Update app.py

app.py

@@ -10,7 +10,7 @@ import urllib.request
 from PIL import Image, ImageDraw
 import matplotlib.pyplot as plt
 
-from Garage.models.GroundedSegmentAnything.segment_anything.segment_anything import SamPredictor, sam_model_registry
+from Garage.models.GroundedSegmentAnything.segment_anything.segment_anything import SamPredictor, build_sam
 from Garage.models.GroundedSegmentAnything.GroundingDINO.groundingdino.util.inference import Model
 from Garage import Augmenter
 
@@ -22,7 +22,7 @@ MODEL_DICT = dict(
 )
 
 GROUNDING_DINO_CONFIG_PATH = "Garage/models/GroundedSegmentAnything/GroundingDINO_SwinT_OGC.py"
-GROUNDING_DINO_CHECKPOINT_PATH = "Garage/models/checkpoints/GroundedSegmentAnything/groundingdino_swint_ogc.pth"
+GROUNDING_DINO_CHECKPOINT_PATH = "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth"
 SAM_CHECKPOINT_PATH = "Garage/models/checkpoints/GroundedSegmentAnything/sam_vit_h_4b8939.pth"
 SAM_ENCODER_VERSION = "vit_h"
 
@@ -210,7 +210,7 @@ class GradioWindow():
                 print(f"Model {model} already exists")
 
     def setup_model(self) -> SamPredictor:
-        self.sam = sam_model_registry[self.model_type](checkpoint=self.SAM_CHECKPOINT_PATH)
+        self.sam = build_sam(checkpoint='sam_vit_h_4b8939.pth')
         self.sam.to(device=self.device)
         self.sam_predictor = SamPredictor(self.sam)
 
@@ -219,6 +219,8 @@ class GradioWindow():
             model_checkpoint_path=self.GROUNDING_DINO_CHECKPOINT_PATH, 
             device=self.device
             )
+        
+        print("MODELS LOADED!")
 
     def change_mask_type(self, image, is_segmmask):
         self.selected_mask = None