Updated SAM notebook (#489)

opengeos · Jul 8, 2023 · 839a2f1 · 839a2f1
1 parent 1eeeddf
commit 839a2f1
Show file tree

Hide file tree

Showing 2 changed files with 128 additions and 74 deletions.
diff --git a/docs/notebooks/75_segment_anything.ipynb b/docs/notebooks/75_segment_anything.ipynb
@@ -29,15 +29,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "outputId": "6924c5c8-39a0-4ac6-f114-9f1f8d102e88",
-    "vscode": {
-     "languageId": "python"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# %pip install leafmap"
@@ -67,8 +59,7 @@
    "source": [
     "import os\n",
     "import leafmap\n",
-    "import torch\n",
-    "from samgeo import SamGeo, tms_to_geotiff"
+    "from samgeo import SamGeo, tms_to_geotiff, get_basemaps"
    ]
   },
   {
@@ -85,7 +76,7 @@
    "outputs": [],
    "source": [
     "m = leafmap.Map(center=[29.676840, -95.369222], zoom=19)\n",
-    "m.add_basemap('SATELLITE')\n",
+    "m.add_basemap(\"SATELLITE\")\n",
     "m"
    ]
   },
@@ -123,8 +114,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "image = 'satellite.tif'\n",
-    "tms_to_geotiff(output=image, bbox=bbox, zoom=20, source='Satellite')"
+    "image = \"satellite.tif\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Besides the `satellite` basemap, you can use any of the following basemaps returned by the `get_basemaps()` function:"
    ]
   },
   {
@@ -133,7 +130,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# image = '/path/to/your/own/image.tif'"
+    "# get_basemaps().keys()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Specify the basemap as the source."
    ]
   },
   {
@@ -142,22 +146,30 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "m.add_raster(image, layer_name='Image')\n",
-    "m"
+    "tms_to_geotiff(output=image, bbox=bbox, zoom=20, source=\"Satellite\", overwrite=True)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "![](https://i.imgur.com/KAm84IY.png)"
+    "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": [
-    "## Initialize SAM class"
+    "Display the downloaded image on the map."
    ]
   },
   {
@@ -166,8 +178,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "out_dir = os.path.join(os.path.expanduser('~'), 'Downloads')\n",
-    "checkpoint = os.path.join(out_dir, 'sam_vit_h_4b8939.pth')"
+    "m.layers[-1].visible = False  # turn off the basemap\n",
+    "m.add_raster(image, layer_name=\"Image\")\n",
+    "m"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "![](https://i.imgur.com/KAm84IY.png)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Initialize SAM class"
    ]
   },
   {
@@ -176,13 +203,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
     "sam = SamGeo(\n",
-    "    checkpoint=checkpoint,\n",
-    "    model_type='vit_h',\n",
-    "    device=device,\n",
-    "    erosion_kernel=(3, 3),\n",
-    "    mask_multiplier=255,\n",
+    "    model_type=\"vit_h\",\n",
+    "    checkpoint=\"sam_vit_h_4b8939.pth\",\n",
     "    sam_kwargs=None,\n",
     ")"
    ]
@@ -191,7 +214,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Segment the image"
+    "## Segment the image\n",
+    "\n",
+    "Set `batch=True` to segment the image in batches. This is useful for large images that cannot fit in memory."
    ]
   },
   {
@@ -200,8 +225,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "mask = 'segment.tiff'\n",
-    "sam.generate(image, mask)"
+    "mask = \"segment.tif\"\n",
+    "sam.generate(\n",
+    "    image, mask, batch=True, foreground=True, erosion_kernel=(3, 3), mask_multiplier=255\n",
+    ")"
    ]
   },
   {
@@ -219,7 +246,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "vector = 'segment.gpkg'\n",
+    "vector = \"segment.gpkg\"\n",
     "sam.tiff_to_gpkg(mask, vector, simplify_tolerance=None)"
    ]
   },
@@ -236,7 +263,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "shapefile = 'segment.shp'\n",
+    "shapefile = \"segment.shp\"\n",
     "sam.tiff_to_vector(mask, shapefile)"
    ]
   },
@@ -254,12 +281,12 @@
    "outputs": [],
    "source": [
     "style = {\n",
-    "    'color': '#3388ff',\n",
-    "    'weight': 2,\n",
-    "    'fillColor': '#7c4185',\n",
-    "    'fillOpacity': 0.5,\n",
+    "    \"color\": \"#3388ff\",\n",
+    "    \"weight\": 2,\n",
+    "    \"fillColor\": \"#7c4185\",\n",
+    "    \"fillOpacity\": 0.5,\n",
     "}\n",
-    "m.add_vector(vector, layer_name='Vector', style=style)\n",
+    "m.add_vector(vector, layer_name=\"Vector\", style=style)\n",
     "m"
    ]
   },