Update app.py

app.py

@@ -8,6 +8,8 @@ from PIL import Image
 from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 
+
+
 # Remove all CUDA-specific configurations
 torch.autocast(device_type="cpu", dtype=torch.float32).__enter__()
 
@@ -51,7 +53,7 @@ def show_mask(mask, ax, random_color=False, borders=True):
         mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2)
     ax.imshow(mask_image)
 
-def show_points(coords, labels, ax, marker_size=375):
+def show_points(coords, labels, ax, marker_size=200):
     pos_points = coords[labels==1]
     neg_points = coords[labels==0]
     ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
@@ -112,16 +114,18 @@ def process_mask(mask, expand_contract_px, expand, feathering_enabled, feather_s
 def sam_process(input_image, checkpoint, tracking_points, trackings_input_label, expand_contract_px, expand, feathering_enabled, feather_size):
     image = Image.open(input_image)
     image = np.array(image.convert("RGB"))
-    checkpoint_map = {
-        "tiny": ("./checkpoints/sam2_hiera_tiny.pt", "sam2_hiera_t.yaml"),
-        "small": ("./checkpoints/sam2_hiera_small.pt", "sam2_hiera_s.yaml"),
-        "base-plus": ("./checkpoints/sam2_hiera_base_plus.pt", "sam2_hiera_b+.yaml"),
-        "large": ("./checkpoints/sam2_hiera_large.pt", "sam2_hiera_l.yaml")
+    sam21_hfmap = {
+        "tiny": "facebook/sam2.1-hiera-tiny",
+        "small": "facebook/sam2.1-hiera-small",
+        "base-plus": "facebook/sam2.1-hiera-base-plus",
+        "large": "facebook/sam2.1-hiera-large",
     }
-    sam2_checkpoint, model_cfg = checkpoint_map[checkpoint]
+    # sam2_checkpoint, model_cfg = checkpoint_map[checkpoint]
     # Use CPU for both model and computations
-    sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cpu")
-    predictor = SAM2ImagePredictor(sam2_model)
+    # sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cpu")
+    predictor = SAM2ImagePredictor.from_pretrained(sam21_hfmap[checkpoint], device="cpu")
+
+    # predictor = SAM2ImagePredictor(sam2_model)
     predictor.set_image(image)
     input_point = np.array(tracking_points.value)
     input_label = np.array(trackings_input_label.value)