Update README.md

README.md

@@ -23,11 +23,11 @@ The training was performed for 100 epochs with a batch size of 160 using the Ada
 
 ```python
 import requests
-import torch
 import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
-from transformers import SamModel, SamProcessor
+from transformers import SamModel, SamProcessor, SamImageProcessor
+import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -40,7 +40,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:
@@ -62,7 +62,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")