Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add HQ-SAM notebooks #162

Merged
merged 1 commit into from
Aug 15, 2023
Merged

Add HQ-SAM notebooks #162

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
382 changes: 382 additions & 0 deletions docs/examples/automatic_mask_generator_hq.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
}
Loading