Update app.py

app.py

@@ -44,11 +44,14 @@ def segment_image(input_image, selected_model):
     for result in results:
         boxes = result.boxes
         bbox = boxes.xyxy.tolist()[0]
-    
-    sam_checkpoint = "medsam_vit_b.pth"
-    model_type = "vit_b"
-    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
-    sam.to(device='cpu')
+    sam_checkpoint = "sam_vit_h_4b8939.pth"
+    model_type = "vit_h"
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint, map_location=torch.device('cpu'))
+
+    #sam_checkpoint = "medsam_vit_b.pth"
+    #model_type = "vit_b"
+    #sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+    #sam.to(device='cpu')
     predictor = SamPredictor(sam)
     predictor.set_image(img)