test SAM 1

app.py

@@ -28,7 +28,8 @@ from threeDFixer.datasets.utils import (
     project2ply
 )
 from threeDFixer.utils import render_utils, postprocessing_utils
-from scripts.grounding_sam2 import plot_segmentation, segment
+from transformers import AutoModelForMaskGeneration, AutoProcessor
+from scripts.grounding_sam import plot_segmentation, segment
 from sam2.build_sam import build_sam2
 from sam2.sam2_image_predictor import SAM2ImagePredictor
 import copy
@@ -55,18 +56,11 @@ generated_object_map = {}
 
 # Prepare models
 ## Grounding SAM
-if not os.path.exists("./checkpoints/sam2.1_hiera_large.pt"):
-    os.makedirs("./checkpoints/", exist_ok=True)
-    os.system("wget -O ./checkpoints/sam2.1_hiera_large.pt https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt")
-SAM2_CHECKPOINT = "./checkpoints/sam2.1_hiera_large.pt"
-SAM2_CONFIG = "configs/sam2.1/sam2.1_hiera_l.yaml"
-
-sam2_model = build_sam2(
-    config_file=SAM2_CONFIG,
-    ckpt_path=SAM2_CHECKPOINT,
-    device=DEVICE
+segmenter_id = "facebook/sam-vit-base"
+sam_processor = AutoProcessor.from_pretrained(segmenter_id)
+sam_segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to(
+    DEVICE, DTYPE
 )
-sam2_predictor = SAM2ImagePredictor(sam2_model)
 
 ############## 3D-Fixer model
 model_dir = 'HorizonRobotics/3D-Fixer'
@@ -192,14 +186,6 @@ def run_segmentation(
         rgb_image = image_prompts["image"].convert("RGB")
 
         global work_space
-        # global sam2_predictor
-
-        # if sam2_predictor is None:
-        #     sam2_model = build_sam2(
-        #         config_file=SAM2_CONFIG,
-        #         ckpt_path=SAM2_CHECKPOINT,
-        #     )
-        #     sam2_predictor = SAM2ImagePredictor(sam2_model)
 
         # pre-process the layers and get the xyxy boxes of each layer
         if len(image_prompts["points"]) == 0:
@@ -214,7 +200,8 @@ def run_segmentation(
         ]
 
         detections = segment(
-            sam2_predictor,
+            sam_processor,
+            sam_segmentator,
             rgb_image,
             boxes=[boxes],
             polygon_refinement=polygon_refinement,