Upload 2 files

scripts/pt_example.py

@@ -1,9 +1,9 @@
 import requests
-import torch
 import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
 from transformers import SamModel, SamProcessor
+import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -16,7 +16,7 @@ input_boxes = [95., 255., 190., 350.]
 
 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())
+probs = processor.image_processor.post_process_masks(outputs.pred_masks.sigmoid().cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu(), binarize=False)
 
 def show_mask(mask, ax, random_color):
     if random_color:
@@ -38,7 +38,7 @@ 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_mask(mask=probs[0] > 0.5, ax=ax[1], random_color=False)
 show_box(input_boxes, ax[1])
 ax[1].set_title("MedSAM Segmentation")
 ax[1].axis("off")

scripts/tf_example.py

@@ -3,6 +3,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
 from transformers import TFSamModel, SamProcessor
+import tensorflow as tf
 
 model = TFSamModel.from_pretrained("flaviagiammarino/medsam-vit-base")
 processor = SamProcessor.from_pretrained("flaviagiammarino/medsam-vit-base")
@@ -13,7 +14,7 @@ input_boxes = [95., 255., 190., 350.]
 
 inputs = processor(raw_image, input_boxes=[[input_boxes]], return_tensors="tf")
 outputs = model(**inputs, multimask_output=False)
-masks = processor.image_processor.post_process_masks([outputs.pred_masks.numpy()[0],], inputs["original_sizes"].numpy(), inputs["reshaped_input_sizes"].numpy())
+probs = processor.image_processor.post_process_masks([tf.sigmoid(outputs.pred_masks).numpy()[0],], inputs["original_sizes"].numpy(), inputs["reshaped_input_sizes"].numpy(), binarize=False)
 
 def show_mask(mask, ax, random_color):
     if random_color:
@@ -35,7 +36,7 @@ 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_mask(mask=probs[0] > 0.5, ax=ax[1], random_color=False)
 show_box(input_boxes, ax[1])
 ax[1].set_title("MedSAM Segmentation")
 ax[1].axis("off")