Add SAM 2 notebook examples

opengeos · Sep 16, 2024 · 0353f00 · 0353f00
1 parent 60c27bd
commit 0353f00
Show file tree

Hide file tree

Showing 2 changed files with 264 additions and 0 deletions.
diff --git a/docs/examples/sam2_predictor.ipynb b/docs/examples/sam2_predictor.ipynb
@@ -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
+}
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -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