Update mask_adapter/sam_maskadapter.py

mask_adapter/sam_maskadapter.py

@@ -227,18 +227,18 @@ class SAMPointVisualizationDemo(object):
         self.mask_adapter = mask_adapter
 
         
-        from .data.datasets import openseg_classes
+        #from .data.datasets import openseg_classes
 
-        COCO_CATEGORIES_pan = openseg_classes.get_coco_categories_with_prompt_eng()
+        #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]
+        #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]
+        #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.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() 
@@ -248,9 +248,11 @@ class SAMPointVisualizationDemo(object):
         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]
-        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]
-        return thing_classes + stuff_classes + lvis_classes
+        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
+        return [ class_name  for class_name in class_names ]
 
 
     def extract_features_convnext(self, x):
@@ -280,7 +282,9 @@ class SAMPointVisualizationDemo(object):
         
         return clip_vis_dense
     
-    def run_on_image_with_points(self, ori_image, points,text_features):
+    def run_on_image_with_points(self, ori_image, points,text_features,class_names=None):
+        if class_names != None:
+            self.class_names = class_names
         height, width, _ = ori_image.shape
 
         image = ori_image
@@ -357,7 +361,10 @@ class SAMPointVisualizationDemo(object):
 
         return None, Image.fromarray(overlay)
     
-    def run_on_image_with_boxes(self, ori_image, bbox,text_features):
+    def run_on_image_with_boxes(self, ori_image, bbox,text_features,class_names=None):
+        if class_names != None:
+            self.class_names = class_names
+            
         height, width, _ = ori_image.shape
 
         image = ori_image