Update mask_adapter/sam_maskadapter.py

mask_adapter/sam_maskadapter.py

@@ -131,7 +131,7 @@ class SAMVisualizationDemo(object):
         pred_masks = np.row_stack(pred_masks)
         pred_masks = BitMasks(pred_masks)
 
-        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
+        image = torch.as_tensor(image.astype("float16").transpose(2, 0, 1))
 
         pixel_mean = torch.tensor(PIXEL_MEAN).view(-1, 1, 1)
         pixel_std = torch.tensor(PIXEL_STD).view(-1, 1, 1)
@@ -150,13 +150,13 @@ class SAMVisualizationDemo(object):
             # text_features = self.clip_model.encode_text(text)
             # text_features /= text_features.norm(dim=-1, keepdim=True)
             
-            features = self.extract_features_convnext(image.float())
+            features = self.extract_features_convnext(image))
             
             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())
+            semantic_activation_maps = self.mask_adapter(clip_vis_dense, pred_masks.tensor.unsqueeze(0).half())
             
             maps_for_pooling = F.interpolate(semantic_activation_maps, size=clip_feature.shape[-2:],
                                                 mode='bilinear', align_corners=False)
@@ -295,7 +295,7 @@ class SAMPointVisualizationDemo(object):
 
         pred_masks = BitMasks(masks)
 
-        image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
+        image = torch.as_tensor(image.astype("float16").transpose(2, 0, 1))
 
         pixel_mean = torch.tensor(PIXEL_MEAN).view(-1, 1, 1)
         pixel_std = torch.tensor(PIXEL_STD).view(-1, 1, 1)
@@ -311,12 +311,12 @@ class SAMPointVisualizationDemo(object):
             # 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.float())
+            features = self.extract_features_convnext(image)
             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())
+            semantic_activation_maps = self.mask_adapter(clip_vis_dense, pred_masks.tensor.unsqueeze(0).half())
             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)