Update mask_adapter/sam_maskadapter.py

mask_adapter/sam_maskadapter.py

@@ -81,8 +81,8 @@ class SAMVisualizationDemo(object):
                                                 crop_n_points_downscale_factor=2,
                                                 min_mask_region_area=100)
 
-        self.clip_model = clip_model.to(self.cpu_device)
-        self.mask_adapter = mask_adapter.to(self.cpu_device)
+        self.clip_model = clip_model
+        self.mask_adapter = mask_adapter
         
 
         
@@ -147,17 +147,17 @@ class SAMVisualizationDemo(object):
 
 
         with torch.no_grad():
-            self.clip_model.to(self.cpu_device)
-            text_features = self.clip_model.encode_text(text.to(self.cpu_device))
+            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.to(self.cpu_device).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).float().to(self.cpu_device))
+            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 +188,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().cpu())
+        semseg = torch.einsum("qc,qhw->chw", select_cls.float(), pred_masks.tensor.float().cuda())
 
         r = semseg
         blank_area = (r[0] == 0)
@@ -222,9 +222,9 @@ class SAMPointVisualizationDemo(object):
 
         self.predictor = SAM2ImagePredictor(sam2)
 
-        self.clip_model = clip_model.to(self.cpu_device)
+        self.clip_model = clip_model
 
-        self.mask_adapter = mask_adapter.to(self.cpu_device)
+        self.mask_adapter = mask_adapter
 
         
         from .data.datasets import openseg_classes
@@ -239,7 +239,7 @@ class SAMPointVisualizationDemo(object):
         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")).to("cpu")
+        self.text_embedding = torch.from_numpy(np.load("./text_embedding/lvis_coco_text_embedding.npy")).to("cuda")
     
         self.class_names = self._load_class_names() 
 
@@ -308,17 +308,17 @@ class SAMPointVisualizationDemo(object):
         # text = open_clip.tokenize(txts)
 
         with torch.no_grad():
-            self.clip_model.to(self.cpu_device)
+            self.clip_model.cuda()
             # 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 = self.text_embedding
-            features = self.extract_features_convnext(image.cpu().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).float().cpu())
+            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)
 
             B, C = clip_feature.size(0), clip_feature.size(1)