flaviagiammarino
commited on
Commit
•
0a901f8
1
Parent(s):
b00dede
Update README.md
Browse files
README.md
CHANGED
@@ -1,3 +1,58 @@
|
|
1 |
---
|
2 |
license: apache-2.0
|
3 |
---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
---
|
2 |
license: apache-2.0
|
3 |
---
|
4 |
+
# Model Card for MedSAM
|
5 |
+
|
6 |
+
MedSAM is a fine-tuned version of [SAM](https://huggingface.co/docs/transformers/main/model_doc/sam) for the medical domain.
|
7 |
+
|
8 |
+
## Usage
|
9 |
+
|
10 |
+
```python
|
11 |
+
import requests
|
12 |
+
import torch
|
13 |
+
import numpy as np
|
14 |
+
import matplotlib.pyplot as plt
|
15 |
+
from PIL import Image
|
16 |
+
from transformers import SamModel, SamProcessor
|
17 |
+
|
18 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
19 |
+
|
20 |
+
model = SamModel.from_pretrained("flaviagiammarino/medsam-vit-base").to(device)
|
21 |
+
processor = SamProcessor.from_pretrained("flaviagiammarino/medsam-vit-base")
|
22 |
+
|
23 |
+
img_url = "https://raw.githubusercontent.com/bowang-lab/MedSAM/main/assets/img_demo.png"
|
24 |
+
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
|
25 |
+
input_boxes = [95., 255., 190., 350.]
|
26 |
+
|
27 |
+
inputs = processor(raw_image, input_boxes=[[input_boxes]], return_tensors="pt").to(device)
|
28 |
+
outputs = model(**inputs, multimask_output=False)
|
29 |
+
masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())
|
30 |
+
|
31 |
+
def show_mask(mask, ax, random_color):
|
32 |
+
if random_color:
|
33 |
+
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
|
34 |
+
else:
|
35 |
+
color = np.array([251/255, 252/255, 30/255, 0.6])
|
36 |
+
h, w = mask.shape[-2:]
|
37 |
+
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
|
38 |
+
ax.imshow(mask_image)
|
39 |
+
|
40 |
+
def show_box(box, ax):
|
41 |
+
x0, y0 = box[0], box[1]
|
42 |
+
w, h = box[2] - box[0], box[3] - box[1]
|
43 |
+
ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor="blue", facecolor=(0, 0, 0, 0), lw=2))
|
44 |
+
|
45 |
+
fig, ax = plt.subplots(1, 2, figsize=(10, 5))
|
46 |
+
ax[0].imshow(np.array(raw_image))
|
47 |
+
show_box(input_boxes, ax[0])
|
48 |
+
ax[0].set_title("Input Image and Bounding Box")
|
49 |
+
ax[0].axis("off")
|
50 |
+
ax[1].imshow(np.array(raw_image))
|
51 |
+
show_mask(masks[0], ax=ax[1], random_color=False)
|
52 |
+
show_box(input_boxes, ax[1])
|
53 |
+
ax[1].set_title("MedSAM Segmentation")
|
54 |
+
ax[1].axis("off")
|
55 |
+
plt.tight_layout()
|
56 |
+
plt.show()
|
57 |
+
```
|
58 |
+
![example](scripts/pt_example.png)
|