sumit-ai-ml commited on
Commit
235cdb3
·
1 Parent(s): f944c1e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -7
app.py CHANGED
@@ -44,14 +44,12 @@ def segment_image(input_image, selected_model):
44
  for result in results:
45
  boxes = result.boxes
46
  bbox = boxes.xyxy.tolist()[0]
47
- sam_checkpoint = "medsam_vit_b.pth"
48
- model_type = "vit_b"
49
- sam = sam_model_registry[model_type](checkpoint=sam_checkpoint, map_location=torch.device('cpu'))
50
 
51
- #sam_checkpoint = "medsam_vit_b.pth"
52
- #model_type = "vit_b"
53
- #sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
54
- #sam.to(device='cpu')
55
  predictor = SamPredictor(sam)
56
  predictor.set_image(img)
57
 
 
44
  for result in results:
45
  boxes = result.boxes
46
  bbox = boxes.xyxy.tolist()[0]
47
+
 
 
48
 
49
+ sam_checkpoint = "sam_vit_h_4b8939.pth"
50
+ model_type = "vit_h"
51
+ sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
52
+ sam.to(device='cpu')
53
  predictor = SamPredictor(sam)
54
  predictor.set_image(img)
55