add vis

app/models/detectors/rapsam.py

@@ -60,7 +60,7 @@ class YOSOVideoSam(Mask2formerVideo):
                 iou_results=iou_results,
                 rescale=False
             )
-            mask_pred_results = results_list[0]['pan_results'].sem_seg[None]
-            mask_cls_results = mask_cls_results
+            # mask_pred_results = results_list[0]['pan_results'].sem_seg[None]
+            return results_list, None
 
         return mask_pred_results.cpu().numpy(), mask_cls_results.cpu().numpy()

main.py

@@ -9,9 +9,12 @@ from PIL import Image
 # mm libs
 from mmdet.registry import MODELS
 from mmdet.structures import DetDataSample
+from mmdet.visualization import DetLocalVisualizer
 from mmengine import Config, print_log
 from mmengine.structures import InstanceData
 
+from mmdet.datasets.coco_panoptic import CocoPanopticDataset
+
 from PIL import ImageDraw
 
 IMG_SIZE = 1024
@@ -30,6 +33,7 @@ model.init_weights()
 mean = torch.tensor([123.675, 116.28, 103.53], device=device)[:, None, None]
 std = torch.tensor([58.395, 57.12, 57.375], device=device)[:, None, None]
 
+visualizer = DetLocalVisualizer()
 
 examples = [
     ["assets/000000000139.jpg"],
@@ -126,22 +130,32 @@ def segment_point(image, img_state):
         batch_data_samples[0].gt_instances_collected = gt_instances
         batch_data_samples[0].set_metainfo(dict(batch_input_shape=(im_h, im_w)))
         batch_data_samples[0].set_metainfo(dict(img_shape=(h, w)))
+        is_prompt = True
     else:
         batch_data_samples = [DetDataSample()]
         batch_data_samples[0].set_metainfo(dict(batch_input_shape=(im_h, im_w)))
         batch_data_samples[0].set_metainfo(dict(img_shape=(h, w)))
+        is_prompt = False
     with torch.no_grad():
         masks, cls_pred = model.predict_with_point(img_tensor, batch_data_samples)
 
-    masks = masks[0, 0, :h, :w]
-    masks = masks > 0.
-    rgb_shape = tuple(list(masks.shape) + [3])
-    color = np.zeros(rgb_shape, dtype=np.uint8)
-    color[masks] = np.array([97, 217, 54])
-    # color[masks] = np.array([217, 90, 54])
-    output_img = (output_img * 0.7 + color * 0.3).astype(np.uint8)
-
-    output_img = Image.fromarray(output_img)
+    assert len(masks) == 1
+    masks = masks[0]
+    if is_prompt:
+        masks = masks[0, :h, :w]
+        masks = masks > 0.   # no sigmoid
+        rgb_shape = tuple(list(masks.shape) + [3])
+        color = np.zeros(rgb_shape, dtype=np.uint8)
+        color[masks] = np.array([97, 217, 54])
+        output_img = (output_img * 0.7 + color * 0.3).astype(np.uint8)
+        output_img = Image.fromarray(output_img)
+    else:
+        output_img = visualizer._draw_panoptic_seg(
+            output_img,
+            masks['pan_results'].to('cpu').numpy(),
+            classes=CocoPanopticDataset.METAINFO['classes'],
+            palette=CocoPanopticDataset.METAINFO['palette']
+        )
     return image, output_img