From 952c2eb951f774d96e99e3f3632c8e82fe8b4b69 Mon Sep 17 00:00:00 2001 From: Qiusheng Wu Date: Mon, 14 Aug 2023 23:34:27 -0400 Subject: [PATCH] Add HQ-SAM notebooks --- README.md | 2 +- .../automatic_mask_generator_hq.ipynb | 382 ++++++++++++++++++ docs/examples/input_prompts_hq.ipynb | 266 ++++++++++++ docs/index.md | 2 +- mkdocs.yml | 2 + requirements.txt | 1 + samgeo/hq_sam.py | 19 +- 7 files changed, 656 insertions(+), 18 deletions(-) create mode 100644 docs/examples/automatic_mask_generator_hq.ipynb create mode 100644 docs/examples/input_prompts_hq.ipynb diff --git a/README.md b/README.md index 7e1768fb..482ffce1 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ The **segment-geospatial** package draws its inspiration from [segment-anything- ## Features - Download map tiles from Tile Map Service (TMS) servers and create GeoTIFF files -- Segment GeoTIFF files using the Segment Anything Model (SAM) +- Segment GeoTIFF files using the Segment Anything Model ([SAM](https://github.com/facebookresearch/segment-anything)) and [HQ-SAM](https://github.com/SysCV/sam-hq) - Segment remote sensing imagery with text prompts - Create foreground and background markers interactively - Load existing markers from vector datasets diff --git a/docs/examples/automatic_mask_generator_hq.ipynb b/docs/examples/automatic_mask_generator_hq.ipynb new file mode 100644 index 00000000..a510a0aa --- /dev/null +++ b/docs/examples/automatic_mask_generator_hq.ipynb @@ -0,0 +1,382 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Automatically generating object masks with HQ-SAM\n", + "\n", + "[![image](https://studiolab.sagemaker.aws/studiolab.svg)](https://studiolab.sagemaker.aws/import/github/opengeos/segment-geospatial/blob/main/docs/examples/automatic_mask_generator_hq.ipynb)\n", + "[![image](https://img.shields.io/badge/Open-Planetary%20Computer-black?style=flat&logo=microsoft)](https://pccompute.westeurope.cloudapp.azure.com/compute/hub/user-redirect/git-pull?repo=https://github.com/opengeos/segment-geospatial&urlpath=lab/tree/segment-geospatial/docs/examples/automatic_mask_generator_hq.ipynb&branch=main)\n", + "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/segment-geospatial/blob/main/docs/examples/automatic_mask_generator_hq.ipynb)\n", + "\n", + "This notebook shows how to segment objects from an image using the High-Quality Segment Anything Model ([HQ-SAM](https://github.com/SysCV/sam-hq)) with a few lines of code. " + ] + }, + { + "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 segment-geospatial" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import leafmap\n", + "from samgeo.hq_sam import SamGeo, show_image, download_file, overlay_images, tms_to_geotiff" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Create an interactive map" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m = leafmap.Map(center=[37.8713, -122.2580], zoom=17, height=\"800px\")\n", + "m.add_basemap(\"SATELLITE\")\n", + "m" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\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_bounds() is not None:\n", + " bbox = m.user_roi_bounds()\n", + "else:\n", + " bbox = [-122.2659, 37.8682, -122.2521, 37.8741]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Download a sample image" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "image = \"satellite.tif\"\n", + "tms_to_geotiff(output=image, bbox=bbox, zoom=17, source=\"Satellite\", overwrite=True)" + ] + }, + { + "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\n", + "\n", + "Specify the file path to the model checkpoint. If it is not specified, the model will to downloaded to the working directory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam = SamGeo(\n", + " model_type=\"vit_h\", # can be vit_h, vit_b, vit_l, vit_tiny\n", + " sam_kwargs=None,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Automatic mask generation\n", + "\n", + "Segment the image and save the results to a GeoTIFF file. Set `unique=True` to assign a unique ID to each object. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam.generate(image, output=\"masks.tif\", foreground=True, unique=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam.show_masks(cmap=\"binary_r\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Show the object annotations (objects with random color) on the map." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam.show_anns(axis=\"off\", alpha=1, output=\"annotations.tif\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compare images with a slider." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "leafmap.image_comparison(\n", + " \"satellite.tif\",\n", + " \"annotations.tif\",\n", + " label1=\"Satellite Image\",\n", + " label2=\"Image Segmentation\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Add image to the map." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m.add_raster(\"annotations.tif\", alpha=0.5, layer_name=\"Masks\")\n", + "m" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Convert the object annotations to vector format, such as GeoPackage, Shapefile, or GeoJSON." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam.tiff_to_vector(\"masks.tif\", \"masks.gpkg\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Automatic mask generation options\n", + "\n", + "There are several tunable parameters in automatic mask generation that control how densely points are sampled and what the thresholds are for removing low quality or duplicate masks. Additionally, generation can be automatically run on crops of the image to get improved performance on smaller objects, and post-processing can remove stray pixels and holes. Here is an example configuration that samples more masks:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam_kwargs = {\n", + " \"points_per_side\": 32,\n", + " \"pred_iou_thresh\": 0.86,\n", + " \"stability_score_thresh\": 0.92,\n", + " \"crop_n_layers\": 1,\n", + " \"crop_n_points_downscale_factor\": 2,\n", + " \"min_mask_region_area\": 100,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam = SamGeo(\n", + " model_type=\"vit_h\",\n", + " sam_kwargs=sam_kwargs,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "metadata": {}, + "outputs": [], + "source": [ + "sam.generate(image, output=\"masks2.tif\", foreground=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam.show_masks(cmap=\"binary_r\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam.show_anns(axis=\"off\", opacity=1, output=\"annotations2.tif\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compare images with a slider." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "leafmap.image_comparison(\n", + " image,\n", + " \"annotations.tif\",\n", + " label1=\"Image\",\n", + " label2=\"Image Segmentation\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Overlay the annotations on the image and use the slider to change the opacity interactively." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "overlay_images(image, \"annotations2.tif\", backend=\"TkAgg\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "![](https://i.imgur.com/I1IhDgz.gif)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "ArcGISPro", + "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.9.16" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/examples/input_prompts_hq.ipynb b/docs/examples/input_prompts_hq.ipynb new file mode 100644 index 00000000..11329aa4 --- /dev/null +++ b/docs/examples/input_prompts_hq.ipynb @@ -0,0 +1,266 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Generating object masks from input prompts with HQ-SAM\n", + "\n", + "[![image](https://studiolab.sagemaker.aws/studiolab.svg)](https://studiolab.sagemaker.aws/import/github/opengeos/segment-geospatial/blob/main/docs/examples/input_prompts_hq.ipynb)\n", + "[![image](https://img.shields.io/badge/Open-Planetary%20Computer-black?style=flat&logo=microsoft)](https://pccompute.westeurope.cloudapp.azure.com/compute/hub/user-redirect/git-pull?repo=https://github.com/opengeos/segment-geospatial&urlpath=lab/tree/segment-geospatial/docs/examples/input_prompts_hq.ipynb&branch=main)\n", + "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/segment-geospatial/blob/main/docs/examples/input_prompts_hq.ipynb)\n", + "\n", + "This notebook shows how to generate object masks from input prompts with the High-Quality Segment Anything Model ([HQ-SAM](https://github.com/SysCV/sam-hq)). \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 segment-geospatial" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import leafmap\n", + "from samgeo.hq_sam import SamGeo, tms_to_geotiff" + ] + }, + { + "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, height=\"800px\")\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", + "tms_to_geotiff(output=image, bbox=bbox, zoom=16, source=\"Satellite\", overwrite=True)" + ] + }, + { + "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\n", + "\n", + "Specify the file path to the model checkpoint. If it is not specified, the model will to downloaded to the working directory." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Set `automatic=False` to disable the `SamAutomaticMaskGenerator` and enable the `SamPredictor`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam = SamGeo(\n", + " model_type=\"vit_h\", # can be vit_h, vit_b, vit_l, vit_tiny\n", + " automatic=False,\n", + " sam_kwargs=None,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Specify the image to segment. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sam.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", + "sam.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", + "sam.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 = sam.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.9.16" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/index.md b/docs/index.md index 4dbea994..71e606dc 100644 --- a/docs/index.md +++ b/docs/index.md @@ -26,7 +26,7 @@ The **segment-geospatial** package draws its inspiration from [segment-anything- ## Features - Download map tiles from Tile Map Service (TMS) servers and create GeoTIFF files -- Segment GeoTIFF files using the Segment Anything Model (SAM) +- Segment GeoTIFF files using the Segment Anything Model ([SAM](https://github.com/facebookresearch/segment-anything)) and [HQ-SAM](https://github.com/SysCV/sam-hq) - Segment remote sensing imagery with text prompts - Create foreground and background markers interactively - Load existing markers from vector datasets diff --git a/mkdocs.yml b/mkdocs.yml index 8de82271..9bb098cc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -51,7 +51,9 @@ nav: - Examples: - examples/satellite.ipynb - examples/automatic_mask_generator.ipynb + - examples/automatic_mask_generator_hq.ipynb - examples/input_prompts.ipynb + - examples/input_prompts_hq.ipynb - examples/box_prompts.ipynb - examples/text_prompts.ipynb - examples/text_prompts_batch.ipynb diff --git a/requirements.txt b/requirements.txt index 47d20acf..4a635aa5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ segment-anything-py +segment-anything-hq opencv-python pycocotools matplotlib diff --git a/samgeo/hq_sam.py b/samgeo/hq_sam.py index d9f6d7ad..a0bfce68 100644 --- a/samgeo/hq_sam.py +++ b/samgeo/hq_sam.py @@ -1,26 +1,13 @@ """ -The source code is adapted from https://github.com/aliaksandr960/segment-anything-eo. Credit to the author Aliaksandr Hancharenka. +Segment remote sensing imagery with HQ-SAM (High Quality Segment Anything Model). +See https://github.com/SysCV/sam-hq """ import os import cv2 import torch import numpy as np - -try: - from segment_anything_hq import ( - sam_model_registry, - SamAutomaticMaskGenerator, - SamPredictor, - ) -except ImportError: - print("Warning: HQ-SAM not available. Installing.") - install_package("segment-anything-hq") - from segment_anything_hq import ( - sam_model_registry, - SamAutomaticMaskGenerator, - SamPredictor, - ) +from segment_anything_hq import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor from .common import *