Skip to content

Commit

Permalink
Add SAM 2 notebook examples
Browse files Browse the repository at this point in the history
  • Loading branch information
giswqs committed Sep 16, 2024
1 parent 60c27bd commit 0353f00
Show file tree
Hide file tree
Showing 2 changed files with 264 additions and 0 deletions.
262 changes: 262 additions & 0 deletions docs/examples/sam2_predictor.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,262 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Generating object masks from input prompts with SAM 2\n",
"\n",
"[![image](https://studiolab.sagemaker.aws/studiolab.svg)](https://studiolab.sagemaker.aws/import/github/opengeos/segment-geospatial/blob/main/docs/examples/sam2_predictor.ipynb)\n",
"[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/segment-geospatial/blob/main/docs/examples/sam2_predictor.ipynb)\n",
"\n",
"This notebook shows how to generate object masks from input prompts with the Segment Anything Model 2 (SAM 2). \n",
"\n",
"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. "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Install dependencies\n",
"\n",
"Uncomment and run the following cell to install the required dependencies."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# %pip install -U segment-geospatial"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import leafmap\n",
"from samgeo import SamGeo2"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Create an interactive map"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"m = leafmap.Map(center=[37.6412, -122.1353], zoom=15)\n",
"m.add_basemap(\"SATELLITE\")\n",
"m"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Download a sample image\n",
"\n",
"Pan and zoom the map to select the area of interest. Use the draw tools to draw a polygon or rectangle on the map"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"if m.user_roi is not None:\n",
" bbox = m.user_roi_bounds()\n",
"else:\n",
" bbox = [-122.1497, 37.6311, -122.1203, 37.6458]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"image = \"satellite.tif\"\n",
"leafmap.map_tiles_to_geotiff(\n",
" output=image, bbox=bbox, zoom=16, source=\"Satellite\", overwrite=True\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can also use your own image. Uncomment and run the following cell to use your own image."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# image = '/path/to/your/own/image.tif'"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Display the downloaded image on the map."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"m.layers[-1].visible = False\n",
"m.add_raster(image, layer_name=\"Image\")\n",
"m"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Initialize SAM class"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Set `automatic=False` to use the predictor mode rather than the automatic mode."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"sam2 = SamGeo2(\n",
" model_id=\"sam2-hiera-large\",\n",
" automatic=False,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Specify the image to segment. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"sam2.set_image(image)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Image segmentation with input points\n",
"\n",
"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.\n",
"\n",
"Try a single point input:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"point_coords = [[-122.1419, 37.6383]]\n",
"sam2.predict(point_coords, point_labels=1, point_crs=\"EPSG:4326\", output=\"mask1.tif\")\n",
"m.add_raster(\"mask1.tif\", layer_name=\"Mask1\", nodata=0, cmap=\"Blues\", opacity=1)\n",
"m"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Try multiple points input:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"point_coords = [[-122.1464, 37.6431], [-122.1449, 37.6415], [-122.1451, 37.6395]]\n",
"sam2.predict(point_coords, point_labels=1, point_crs=\"EPSG:4326\", output=\"mask2.tif\")\n",
"m.add_raster(\"mask2.tif\", layer_name=\"Mask2\", nodata=0, cmap=\"Greens\", opacity=1)\n",
"m"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Interactive segmentation\n",
"\n",
"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."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"m = sam2.show_map()\n",
"m"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"![](https://i.imgur.com/2Nyg9uW.gif)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "sam",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.8"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
2 changes: 2 additions & 0 deletions mkdocs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,8 @@ nav:
- examples/arcgis.ipynb
- examples/maxar_open_data.ipynb
- examples/sam2_automatic.ipynb
- examples/sam2_predictor.ipynb
- examples/sam2_video.ipynb
- Workshops:
- workshops/purdue.ipynb
- workshops/cn_workshop.ipynb
Expand Down

0 comments on commit 0353f00

Please sign in to comment.