| 1281 | |
| 1282 | |
| 1283 | class ObjectSegmenting: |
| 1284 | template_model = True # Add this line to show this is a template model. |
| 1285 | def __init__(self, Text2Box:Text2Box, Segmenting:Segmenting): |
| 1286 | # self.llm = OpenAI(temperature=0) |
| 1287 | self.grounding = Text2Box |
| 1288 | self.sam = Segmenting |
| 1289 | |
| 1290 | |
| 1291 | @prompts(name="Segment the given object", |
| 1292 | description="useful when you only want to segment the certain objects in the picture" |
| 1293 | "according to the given text" |
| 1294 | "like: segment the cat," |
| 1295 | "or can you segment an obeject for me" |
| 1296 | "The input to this tool should be a comma separated string of two, " |
| 1297 | "representing the image_path, the text description of the object to be found") |
| 1298 | def inference(self, inputs): |
| 1299 | image_path, det_prompt = inputs.split(",") |
| 1300 | print(f"image_path={image_path}, text_prompt={det_prompt}") |
| 1301 | image_pil, image = self.grounding.load_image(image_path) |
| 1302 | |
| 1303 | boxes_filt, pred_phrases = self.grounding.get_grounding_boxes(image, det_prompt) |
| 1304 | updated_image_path = self.sam.segment_image_with_boxes(image_pil,image_path,boxes_filt,pred_phrases) |
| 1305 | print( |
| 1306 | f"\nProcessed ObejectSegmenting, Input Image: {image_path}, Object to be Segment {det_prompt}, " |
| 1307 | f"Output Image: {updated_image_path}") |
| 1308 | return updated_image_path |
| 1309 | |
| 1310 | def merge_masks(self, masks): |
| 1311 | ''' |
| 1312 | Args: |
| 1313 | mask (numpy.ndarray): shape N x 1 x H x W |
| 1314 | Outputs: |
| 1315 | new_mask (numpy.ndarray): shape H x W |
| 1316 | ''' |
| 1317 | if type(masks) == torch.Tensor: |
| 1318 | x = masks |
| 1319 | elif type(masks) == np.ndarray: |
| 1320 | x = torch.tensor(masks,dtype=int) |
| 1321 | else: |
| 1322 | raise TypeError("the type of the input masks must be numpy.ndarray or torch.tensor") |
| 1323 | x = x.squeeze(dim=1) |
| 1324 | value, _ = x.max(dim=0) |
| 1325 | new_mask = value.cpu().numpy() |
| 1326 | new_mask.astype(np.uint8) |
| 1327 | return new_mask |
| 1328 | |
| 1329 | def get_mask(self, image_path, text_prompt): |
| 1330 | |
| 1331 | print(f"image_path={image_path}, text_prompt={text_prompt}") |
| 1332 | # image_pil (PIL.Image.Image) -> size: W x H |
| 1333 | # image (numpy.ndarray) -> H x W x 3 |
| 1334 | image_pil, image = self.grounding.load_image(image_path) |
| 1335 | |
| 1336 | boxes_filt, pred_phrases = self.grounding.get_grounding_boxes(image, text_prompt) |
| 1337 | image = cv2.imread(image_path) |
| 1338 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
| 1339 | self.sam.sam_predictor.set_image(image) |
| 1340 | |