Update mask_adapter/sam_maskadapter.py

mask_adapter/sam_maskadapter.py

@@ -132,7 +132,7 @@ class SAMVisualizationDemo(object):
         
         return clip_vis_dense
     
-    def run_on_image(self, ori_image, class_names):
+    def run_on_image(self, ori_image, class_names, text_features):
         height, width, _ = ori_image.shape
         if width > height:
             new_width = 896
@@ -158,25 +158,25 @@ class SAMVisualizationDemo(object):
         image = (image - pixel_mean) / pixel_std
 
         image = image.unsqueeze(0)
-        
-        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)
+
+        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)
 
 
         with torch.no_grad():
-            self.clip_model.cuda()
-            text_features = self.clip_model.encode_text(text.cuda())
-            text_features /= text_features.norm(dim=-1, keepdim=True)
+            # text_features = self.clip_model.encode_text(text)
+            # text_features /= text_features.norm(dim=-1, keepdim=True)
             
-            features = self.extract_features_convnext(image.cuda().float())
+            features = self.extract_features_convnext(image.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).float().cuda())
+            semantic_activation_maps = self.mask_adapter(clip_vis_dense, pred_masks.tensor.unsqueeze(0).to(text_features).float())
             
             maps_for_pooling = F.interpolate(semantic_activation_maps, size=clip_feature.shape[-2:],
                                                 mode='bilinear', align_corners=False)
@@ -207,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.float().cuda())
+        semseg = torch.einsum("qc,qhw->chw", select_cls.float(), pred_masks.tensor.to(text_features).float())
 
         r = semseg
         blank_area = (r[0] == 0)
@@ -244,6 +244,21 @@ 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() 
 
@@ -411,28 +426,12 @@ class SAMPointVisualizationDemo(object):
         image = image.unsqueeze(0)
 
         # txts = [f'a photo of {cls_name}' for cls_name in self.class_names]
-
-                    
-                    
-            # txts.append(f'a photo of {cls_name}')
-        #assert len(self.class_names) * 14 == len(txts)
-        #text = open_clip.tokenize(txts)
+        # text = open_clip.tokenize(txts)
         
         with torch.no_grad():
-            # 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)
-            # text_features = text_features.reshape(text_features.shape[0] // len(VILD_PROMPT), len(VILD_PROMPT), text_features.shape[-1]).mean(1)
+            # text_features = self.clip_model.encode_text(text.cuda())
             # 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())
+            #np.save("/home/yongkangli/Mask-Adapter/text_embedding/lvis_coco_text_embedding.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']