Update mask_adapter/sam_maskadapter.py

mask_adapter/sam_maskadapter.py

@@ -19,6 +19,25 @@ from PIL import Image
 PIXEL_MEAN = [122.7709383, 116.7460125, 104.09373615]
 PIXEL_STD = [68.5005327, 66.6321579, 70.32316305]
 
+
+VILD_PROMPT = [
+    "a photo of a {}.",
+    "This is a photo of a {}",
+    "There is a {} in the scene",
+    "There is the {} in the scene",
+    "a photo of a {} in the scene",
+    "a photo of a small {}.",
+    "a photo of a medium {}.",
+    "a photo of a large {}.",
+    "This is a photo of a small {}.",
+    "This is a photo of a medium {}.",
+    "This is a photo of a large {}.",
+    "There is a small {} in the scene.",
+    "There is a medium {} in the scene.",
+    "There is a large {} in the scene.",
+]
+
+
 class OpenVocabVisualizer(Visualizer):
     def __init__(self, img_rgb, metadata=None, scale=1.0, instance_mode=ColorMode.IMAGE, class_names=None):
         super().__init__(img_rgb, metadata, scale, instance_mode)
@@ -113,7 +132,7 @@ class SAMVisualizationDemo(object):
         
         return clip_vis_dense
     
-    def run_on_image(self, ori_image, class_names, text_features):
+    def run_on_image(self, ori_image, class_names):
         height, width, _ = ori_image.shape
         if width > height:
             new_width = 896
@@ -139,25 +158,25 @@ class SAMVisualizationDemo(object):
         image = (image - pixel_mean) / pixel_std
 
         image = image.unsqueeze(0)
-
-        image = image.to(text_features)
-        # if len(class_names) == 1:
-        #     class_names.append('others')
-        # txts = [f'a photo of {cls_name}' for cls_name in class_names]
-        # text = open_clip.tokenize(txts)
+        
+        if len(class_names) == 1:
+            class_names.append('others')
+        txts = [f'a photo of {cls_name}' for cls_name in class_names]
+        text = open_clip.tokenize(txts)
 
 
         with torch.no_grad():
-            # text_features = self.clip_model.encode_text(text)
-            # text_features /= text_features.norm(dim=-1, keepdim=True)
+            self.clip_model.cuda()
+            text_features = self.clip_model.encode_text(text.cuda())
+            text_features /= text_features.norm(dim=-1, keepdim=True)
             
-            features = self.extract_features_convnext(image.float())
+            features = self.extract_features_convnext(image.cuda().float())
             
             clip_feature = features['clip_vis_dense']
             
             clip_vis_dense = self.visual_prediction_forward_convnext_2d(clip_feature)
             
-            semantic_activation_maps = self.mask_adapter(clip_vis_dense, pred_masks.tensor.unsqueeze(0).to(text_features).float())
+            semantic_activation_maps = self.mask_adapter(clip_vis_dense, pred_masks.tensor.unsqueeze(0).float().cuda())
             
             maps_for_pooling = F.interpolate(semantic_activation_maps, size=clip_feature.shape[-2:],
                                                 mode='bilinear', align_corners=False)
@@ -188,7 +207,7 @@ class SAMVisualizationDemo(object):
                 select_mask.extend(locs[0].tolist())
         for idx in select_mask:
             select_cls[idx] = class_preds[idx]
-        semseg = torch.einsum("qc,qhw->chw", select_cls.float(), pred_masks.tensor.to(text_features).float())
+        semseg = torch.einsum("qc,qhw->chw", select_cls.float(), pred_masks.tensor.float().cuda())
 
         r = semseg
         blank_area = (r[0] == 0)
@@ -225,33 +244,16 @@ class SAMPointVisualizationDemo(object):
         self.clip_model = clip_model
 
         self.mask_adapter = mask_adapter
-
-        
-        #from .data.datasets import openseg_classes
-
-        #COCO_CATEGORIES_pan = openseg_classes.get_coco_categories_with_prompt_eng()
-        #COCO_CATEGORIES_seg = openseg_classes.get_coco_stuff_categories_with_prompt_eng()
-
-        #thing_classes = [k["name"] for k in COCO_CATEGORIES_pan if k["isthing"] == 1]
-        #stuff_classes = [k["name"] for k in COCO_CATEGORIES_pan]
-        #print(coco_metadata)
-        #lvis_classes = open("./mask_adapter/data/datasets/lvis_1203_with_prompt_eng.txt", 'r').read().splitlines()
-        #lvis_classes = [x[x.find(':')+1:] for x in lvis_classes]
-                
-        #self.class_names = thing_classes + stuff_classes + lvis_classes
-        #self.text_embedding = torch.from_numpy(np.load("./text_embedding/lvis_coco_text_embedding.npy"))
     
         self.class_names = self._load_class_names() 
 
     def _load_class_names(self):
         from .data.datasets import openseg_classes
         COCO_CATEGORIES_pan = openseg_classes.get_coco_categories_with_prompt_eng()
-        thing_classes = [k["name"] for k in COCO_CATEGORIES_pan if k["isthing"] == 1]
         stuff_classes = [k["name"] for k in COCO_CATEGORIES_pan]
         ADE20K_150_CATEGORIES_ = openseg_classes.get_ade20k_categories_with_prompt_eng()
-        ade20k_thing_classes = [k["name"] for k in ADE20K_150_CATEGORIES_ if k["isthing"] == 1]
         ade20k_stuff_classes = [k["name"] for k in ADE20K_150_CATEGORIES_]
-        class_names = thing_classes + stuff_classes + ade20k_thing_classes+ ade20k_stuff_classes
+        class_names =  stuff_classes + ade20k_stuff_classes #+ lvis_classes
         return [ class_name  for class_name in class_names ]
 
 
@@ -285,6 +287,12 @@ class SAMPointVisualizationDemo(object):
     def run_on_image_with_points(self, ori_image, points,text_features,class_names=None):
         if class_names != None:
             self.class_names = class_names
+        else:
+            num_templates = []
+            for cls_name in self.class_names:
+                cls_name = cls_name.replace(', ', ',').split(',')#[0]
+                num_templates.append(len(cls_name))
+
         height, width, _ = ori_image.shape
 
         image = ori_image
@@ -333,7 +341,18 @@ class SAMPointVisualizationDemo(object):
             pooled_clip_feature = (pooled_clip_feature.reshape(B, num_instances, 16, -1).mean(dim=-2).contiguous())
 
             class_preds = (100.0 * pooled_clip_feature @ text_features.T).softmax(dim=-1)
-        class_preds = class_preds.squeeze(0)
+            
+        if class_names is None:
+            final_class_preds = []
+            cur_idx = 0
+            for num_t in num_templates: 
+                final_class_preds.append(class_preds[:, :, cur_idx: cur_idx + num_t].max(-1).values)
+                cur_idx += num_t
+            final_class_preds = torch.stack(final_class_preds, dim=-1)
+
+            class_preds = final_class_preds.squeeze(0)
+        else:
+            class_preds = class_preds.squeeze(0)
 
         # Resize mask to match original image size
         pred_mask = cv2.resize(masks.squeeze(0), (width, height), interpolation=cv2.INTER_NEAREST)  # Resize mask to match original image size
@@ -364,7 +383,12 @@ class SAMPointVisualizationDemo(object):
     def run_on_image_with_boxes(self, ori_image, bbox,text_features,class_names=None):
         if class_names != None:
             self.class_names = class_names
-            
+        else:
+            num_templates = []
+            for cls_name in self.class_names:
+                cls_name = cls_name.replace(', ', ',').split(',')#[0]
+                num_templates.append(len(cls_name))
+
         height, width, _ = ori_image.shape
 
         image = ori_image
@@ -387,12 +411,28 @@ class SAMPointVisualizationDemo(object):
         image = image.unsqueeze(0)
 
         # txts = [f'a photo of {cls_name}' for cls_name in self.class_names]
-        # text = open_clip.tokenize(txts)
+
+                    
+                    
+            # txts.append(f'a photo of {cls_name}')
+        #assert len(self.class_names) * 14 == len(txts)
+        #text = open_clip.tokenize(txts)
         
         with torch.no_grad():
-            # text_features = self.clip_model.encode_text(text.cuda())
+            # text_features = []
+            # bs = 128
+            # for idx in range(0, len(text), bs):
+            #     text_features.append(
+            #         self.clip_model.encode_text(text[idx:idx+bs].cuda())
+            #     )
+            # text_features = torch.cat(text_features, dim=0)
+
+            # #text_features = self.clip_model.encode_text(text.cuda())
             # text_features /= text_features.norm(dim=-1, keepdim=True)
-            #np.save("/home/yongkangli/Mask-Adapter/text_embedding/lvis_coco_text_embedding.npy", text_features.cpu().numpy())
+            # text_features = text_features.reshape(text_features.shape[0] // len(VILD_PROMPT), len(VILD_PROMPT), text_features.shape[-1]).mean(1)
+            # text_features /= text_features.norm(dim=-1, keepdim=True)
+            # print(text_features.shape)
+            # np.save("/data/yongkangli/Mask-Adapter-Demo/text_embedding/coco_ade20k_text_embedding_new.npy", text_features.cpu().numpy())
             #text_features = self.text_embedding.to(self.mask_adapter.device)
             features = self.extract_features_convnext(image.to(text_features).float())
             clip_feature = features['clip_vis_dense']
@@ -411,7 +451,18 @@ class SAMPointVisualizationDemo(object):
             pooled_clip_feature = (pooled_clip_feature.reshape(B, num_instances, 16, -1).mean(dim=-2).contiguous())
 
             class_preds = (100.0 * pooled_clip_feature @ text_features.T).softmax(dim=-1)
-        class_preds = class_preds.squeeze(0)
+            
+        if class_names is None:
+            final_class_preds = []
+            cur_idx = 0
+            for num_t in num_templates: 
+                final_class_preds.append(class_preds[:, :, cur_idx: cur_idx + num_t].max(-1).values)
+                cur_idx += num_t
+            final_class_preds = torch.stack(final_class_preds, dim=-1)
+
+            class_preds = final_class_preds.squeeze(0)
+        else:
+            class_preds = class_preds.squeeze(0)
 
         # Resize mask to match original image size
         pred_mask = cv2.resize(masks.squeeze(0), (width, height), interpolation=cv2.INTER_NEAREST)  # Resize mask to match original image size