Spaces:

sumit-ai-ml
/

SAM-with-YOLO

Runtime error

sumit-ai-ml commited on Oct 3, 2023

Commit

f944c1e

1 Parent(s): 057f3bb

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -44,8 +44,8 @@ def segment_image(input_image, selected_model):
     for result in results:
         boxes = result.boxes
         bbox = boxes.xyxy.tolist()[0]
-    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"

     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, map_location=torch.device('cpu'))
     #sam_checkpoint = "medsam_vit_b.pth"