Upload 2 files

scripts/pt_example.py

@@ -7,14 +7,14 @@ from transformers import SamModel, SamProcessor
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-model = SamModel.from_pretrained("flaviagiammarino/medsam-vit-base")
+model = SamModel.from_pretrained("flaviagiammarino/medsam-vit-base").to(device)
 processor = SamProcessor.from_pretrained("flaviagiammarino/medsam-vit-base")
 
 img_url = "https://raw.githubusercontent.com/bowang-lab/MedSAM/main/assets/img_demo.png"
 raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
-input_boxes = [95, 255, 190, 350]
+input_boxes = [95., 255., 190., 350.]
 
-inputs = processor(raw_image, input_boxes=[input_boxes], return_tensors="pt").to(device)
+inputs = processor(raw_image, input_boxes=[[input_boxes]], return_tensors="pt").to(device)
 outputs = model(**inputs, multimask_output=False)
 masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
 
@@ -32,11 +32,15 @@ def show_box(box, ax):
     w, h = box[2] - box[0], box[3] - box[1]
     ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor="blue", facecolor=(0, 0, 0, 0), lw=2))
 
-plt.imshow(np.array(raw_image))
-ax = plt.gca()
-for mask in masks:
-    show_mask(mask, ax=ax, random_color=False)
-show_box(input_boxes, ax)
-plt.axis("off")
+fig, ax = plt.subplots(1, 2, figsize=(10, 5))
+ax[0].imshow(np.array(raw_image))
+show_box(input_boxes, ax[0])
+ax[0].set_title("Input Image and Bounding Box")
+ax[0].axis("off")
+ax[1].imshow(np.array(raw_image))
+show_mask(masks[0], ax=ax[1], random_color=False)
+show_box(input_boxes, ax[1])
+ax[1].set_title("MedSAM Segmentation")
+ax[1].axis("off")
 plt.tight_layout()
 plt.show()