Update backend.py

backend.py

@@ -155,33 +155,45 @@ CONF_THRESHOLD = 0.2
 # -----------------------------
 # 4) הכנה ל-SAM2
 # -----------------------------
-try:
-    from sam2.build_sam import build_sam2
-    from sam2.sam2_image_predictor import SAM2ImagePredictor
 
-    SAM2_CHECKPOINT = "../../checkpoints/sam2.1_hiera_tiny.pt"
-    SAM2_CONFIG = "../../configs/sam2.1/sam2.1_hiera_t.yaml"
+from typing import Any
+import supervision as sv
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+SAM2_CHECKPOINT = "checkpoints/sam2_hiera_small.pt"
+SAM2_CONFIG = "sam2_hiera_s.yaml"
 
-    if not os.path.exists(SAM2_CHECKPOINT):
+def load_sam_image_model(
+    device: torch.device,
+    config: str = SAM2_CONFIG,
+    checkpoint: str = SAM2_CHECKPOINT
+) -> SAM2ImagePredictor:
+    if not os.path.exists(checkpoint):
         print("[SAM2] מודל SAM2 לא נמצא. מנסה להוריד את המודל...")
-        sam2_url = "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt"
-        os.makedirs(os.path.dirname(SAM2_CHECKPOINT), exist_ok=True)
+        sam2_url = "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2_hiera_small.pt"
+        os.makedirs(os.path.dirname(checkpoint), exist_ok=True)
         response = requests.get(sam2_url)
-        with open(SAM2_CHECKPOINT, 'wb') as f:
+        with open(checkpoint, 'wb') as f:
             f.write(response.content)
         print("[SAM2] מודל SAM2 הורד בהצלחה.")
-
-    if not os.path.exists(SAM2_CONFIG):
+    
+    if not os.path.exists(config):
         print("[SAM2] קובץ הקונפיג SAM2 לא נמצא. מנסה להוריד את הקונפיג...")
-        # כאן יש לוודא שקובץ הקונפיג זמין, או לכלול אותו במאגר שלך
-        raise FileNotFoundError(f"לא נמצא קובץ קונפיג SAM2 בנתיב: {SAM2_CONFIG}")
+        sam2_config_url = "https://path_to_your_config/sam2_hiera_s.yaml"  # עדכן את הקישור המתאים
+        os.makedirs(os.path.dirname(config), exist_ok=True)
+        response = requests.get(sam2_config_url)
+        with open(config, 'wb') as f:
+            f.write(response.content)
+        print("[SAM2] קובץ הקונפיג SAM2 הורד בהצלחה.")
 
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = build_sam2(config, checkpoint, device=device)
+    return SAM2ImagePredictor(sam_model=model)
 
-    sam2_model = build_sam2(SAM2_CONFIG, SAM2_CHECKPOINT, device=device)
-    sam2_predictor = SAM2ImagePredictor(sam2_model)
+try:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    sam2_predictor = load_sam_image_model(device=device)
     sam2_predictor.model.to(device)
-
 except Exception as e:
     print(f"[SAM2] לא מצליח לטעון את SAM2: {e}")
     sam2_predictor = None