download sam from huggingface

blurnonymize.py

@@ -7,7 +7,6 @@ import matplotlib.patches as patches
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
-from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 from utils import *
 
@@ -160,38 +159,8 @@ class ImageBlurnonymizer:
         if self.predictor is not None and not force:
             return
 
-        self.download_weights()
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        sam = build_sam2(f"file://{self.model_cfg_name}", f"file://{self.checkpoint_name}", device=self.device)
-        self.predictor = SAM2ImagePredictor(sam)
-
-    def download_weights(self):
-        # TODO: check whether these files already exist, if not, download them
-        # files names are in self.checkpoint_name and self.model_cfg_name
-
-        print(f"Current working directory: {os.getcwd()}")
-        print(f"Files in current dir: {os.listdir()}")
-        print(f"Files in root dir: {os.listdir('/')}")
-
-        print(f"Files in /data: {os.listdir('/data')}")
-
-        checkpoint_url = "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt"
-        cfg_url = "https://raw.githubusercontent.com/facebookresearch/sam2/refs/heads/main/sam2/configs/sam2.1/sam2.1_hiera_s.yaml"
-
-        config_dir = "/data/configs/sam2.1"
-
-        if not os.path.exists(config_dir):
-            os.makedirs(config_dir)
-
-        if not os.path.exists(self.checkpoint_name):
-            print(
-                f"Downloading checkpoint from {checkpoint_url} to {self.checkpoint_name}"
-            )
-            torch.hub.download_url_to_file(checkpoint_url, self.checkpoint_name)
-
-        if not os.path.exists(self.model_cfg_name):
-            print(f"Downloading config from {cfg_url} to {self.model_cfg_name}")
-            torch.hub.download_url_to_file(cfg_url, self.model_cfg_name)
+        self.predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
 
     @staticmethod
     def _smoothen_mask(mask: np.ndarray):
@@ -276,20 +245,21 @@ class ImageBlurnonymizer:
         # Ensure points are valid coordinates (e.g., non-negative)
         points = [[max(0, p[0]), max(0, p[1])] for p in points]
 
-        self.predictor.set_image(image)
-        masks, scores, _ = self.predictor.predict(
-            box=np.array(bbox),  # Predictor might expect numpy array
-            point_coords=np.array(points),
-            point_labels=np.ones(len(points)),  # Label 1 for inclusion
-            multimask_output=True,
-        )
+        with torch.inference_mode(), torch.autocast(self.device, dtype=torch.bfloat16):
+            self.predictor.set_image(image)
+            masks, scores, _ = self.predictor.predict(
+                box=np.array(bbox),  # Predictor might expect numpy array
+                point_coords=np.array(points),
+                point_labels=np.ones(len(points)),  # Label 1 for inclusion
+                multimask_output=True,
+            )
 
-        # Sort masks by score and select the best one
-        sorted_ind = np.argsort(scores)[::-1]
-        best_mask = masks[sorted_ind[0]]
-        best_score = scores[sorted_ind[0]]
+            # Sort masks by score and select the best one
+            sorted_ind = np.argsort(scores)[::-1]
+            best_mask = masks[sorted_ind[0]]
+            best_score = scores[sorted_ind[0]]
 
-        return self._smoothen_mask(best_mask), best_score
+            return self._smoothen_mask(best_mask), best_score
 
     def censor_image_blur(
         self,