facebookresearch / facebookresearch/sam2
Enhance `image_predictor_example.ipynb` with Interactive Point Addition Using Matplotlib
- Dominant language
- Jupyter Notebook
- Stars
- 19.9k
- Forks
- 2.5k
- PR merge metrics
- No merged PRs in 30d
Description
**Issue:**
The current `image_predictor_example.ipynb` provides a good example of using the image predictor. However, it lacks interactivity, which can enhance user experience and facilitate experimentation.
**Proposed Enhancement:**
Introduce an interactive feature that allows users to add positive and negative points by clicking on the image:
- **Left Click:** Add a positive point.
- **Right Click:** Add a negative point.
This can be achieved with a simple Matplotlib-based script consisting of less than 110 lines of code, eliminating the need for complex third-party tools for simple testing.
**Benefits:**
- **Simplicity:** Easier to understand and modify due to fewer lines of code.
- **User-Friendly:** Enhances interactivity without adding external dependencies.
- **Educational Value:** Helps users learn by directly interacting with the model.
**Implementation:**
I have made a concise script that demonstrates this functionality. You can view the complete code in the following [Gist](https://gist.github.com/future-158/52713022ae08e7d69ddc4727fef4379a).
**Example Code Snippet:**
```python
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
from sam2.sam2_image_predictor import SAM2ImagePredictor
import requests
def load_image(url: str) -> Image.Image:
headers = {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/85.0.4183.83 Safari/537.36"
}
image = Image.open(requests.get(url, stream=True, headers=headers).raw)
return image
%matplotlib widget
predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-tiny")
# predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2-hiera-large")
# load example image
url = "https://images.pexels.com/photos/529782/pexels-photo-529782.jpeg?auto=compress&cs=tinysrgb&w=800"
base_img = load_image(url)
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
predictor.set_image(base_img)
img = np.array(base_img)
fig, ax = plt.subplots()
im = ax.imshow(img)
# Remove ticks
ax.set_xticks([])
ax.set_yticks([])
# Remove tick labels
ax.set_xticklabels([])
ax.set_yticklabels([])
# Remove axis labels
ax.set_xlabel("")
ax.set_ylabel("")
plt.tight_layout()
positive_points = []
negative_points = []
mask = None
def inference() -> np.ndarray:
global mask
point_coords = [*positive_points, *negative_points]
point_labels = [1] * len(positive_points) + [0] * len(negative_points)
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
masks, _, _ = predictor.predict(
point_coords=point_coords,
point_labels=point_labels,
multimask_output=False,
)
mask = masks[0] > 0
blended = Image.blend(
Image.new("RGB", base_img.size, (0, 0, 255)), base_img, alpha=0.5
)
composited = Image.composite(blended, base_img, Image.fromarray(mask > 0)).convert(
"RGB"
)
return np.array(composited)
def on_click(event):
"""
Event handler for mouse click events on the plot.
Parameters:
- event: The mouse event.
"""
if event.inaxes:
x, y = event.xdata, event.ydata
if event.button == 1: # Left mouse button
positive_points.append((x, y))
elif event.button == 3: # Right mouse button
negative_points.append((x, y))
new_rgb = inference()
for p in positive_points:
x, y = p
x, y = int(x), int(y)
new_rgb[y - 5 : y + 5, x - 5 : x + 5] = [0, 255, 0]
for p in negative_points:
x, y = p
x, y = int(x), int(y)
new_rgb[y - 5 : y + 5, x - 5 : x + 5] = [255, 0, 0]
im.set_data(new_rgb)
fig.canvas.draw_idle()
cid = fig.canvas.mpl_connect("button_press_event", on_click)
plt.show()
```
**Conclusion:**
when i first ran demo with my image, i draw image with plotly first (cause it show mouse point coordinates) and manually update postive points and negative points one by one.
i think adding this interactive example can make the image_predictor_example.ipynb more engaging.
I'm happy to contribute this example to the repository or provide further assistance if needed.
Contributor guide
Assessment
This issue has not been assessed yet.