Generating object masks from input prompts with SAM¶
This notebook shows how to generate object masks from input prompts with the Segment Anything Model (SAM).
Make sure you use GPU runtime for this notebook. For Google Colab, go to Runtime -> Change runtime type and select GPU as the hardware accelerator.
The notebook is adapted from segment-anything/notebooks/predictor_example.ipynb, but I have made it much easier to save the segmentation results and visualize them.
Install dependencies¶
Uncomment and run the following cell to install the required dependencies.
# %pip install segment-geospatial
import os
import leafmap
from samgeo import SamGeo, tms_to_geotiff
Create an interactive map¶
m = leafmap.Map(center=[37.6412, -122.1353], zoom=15, height="800px")
m.add_basemap("SATELLITE")
m
Download a sample image¶
Pan and zoom the map to select the area of interest. Use the draw tools to draw a polygon or rectangle on the map
if m.user_roi is not None:
bbox = m.user_roi_bounds()
else:
bbox = [-122.1497, 37.6311, -122.1203, 37.6458]
image = "satellite.tif"
tms_to_geotiff(output=image, bbox=bbox, zoom=16, source="Satellite", overwrite=True)
You can also use your own image. Uncomment and run the following cell to use your own image.
# image = '/path/to/your/own/image.tif'
Display the downloaded image on the map.
m.layers[-1].visible = False
m.add_raster(image, layer_name="Image")
m
Initialize SAM class¶
Specify the file path to the model checkpoint. If it is not specified, the model will to downloaded to the working directory.
Set automatic=False to disable the SamAutomaticMaskGenerator and enable the SamPredictor.
sam = SamGeo(
model_type="vit_h",
automatic=False,
sam_kwargs=None,
)
Specify the image to segment.
sam.set_image(image)
Image segmentation with input points¶
A single point can be used to segment an object. The point can be specified as a tuple of (x, y), such as (col, row) or (lon, lat). The points can also be specified as a file path to a vector dataset. For non (col, row) input points, specify the point_crs parameter, which will automatically transform the points to the image column and row coordinates.
Try a single point input:
point_coords = [[-122.1419, 37.6383]]
sam.predict(point_coords, point_labels=1, point_crs="EPSG:4326", output="mask1.tif")
m.add_raster("mask1.tif", layer_name="Mask1", nodata=0, cmap="Blues", opacity=1)
m
Try multiple points input:
point_coords = [[-122.1464, 37.6431], [-122.1449, 37.6415], [-122.1451, 37.6395]]
sam.predict(point_coords, point_labels=1, point_crs="EPSG:4326", output="mask2.tif")
m.add_raster("mask2.tif", layer_name="Mask2", nodata=0, cmap="Greens", opacity=1)
m
Interactive segmentation¶
Display the interactive map and use the marker tool to draw points on the map. Then click on the Segment button to segment the objects. The results will be added to the map automatically. Click on the Reset button to clear the points and the results.
m = sam.show_map()
m
