Skip to content

Download notebook

Automatically generating object masks with HQ-SAM

image image

This notebook shows how to segment objects from an image using the High-Quality Segment Anything Model (HQ-SAM) with a few lines of code.

Install dependencies

Uncomment and run the following cell to install the required dependencies.

1
# %pip install segment-geospatial

1
2
3
import leafmap
from samgeo.hq_sam import SamGeo
from samgeo.common import overlay_images, tms_to_geotiff

Create an interactive map

1
2
3
m = leafmap.Map(center=[37.8713, -122.2580], zoom=17, height="800px")
m.add_basemap("SATELLITE")
m

Pan and zoom the map to select the area of interest. Use the draw tools to draw a polygon or rectangle on the map

1
2
3
4
if m.user_roi_bounds() is not None:
    bbox = m.user_roi_bounds()
else:
    bbox = [-122.2659, 37.8682, -122.2521, 37.8741]

Download a sample image

1
2
image = "satellite.tif"
tms_to_geotiff(output=image, bbox=bbox, zoom=17, source="Satellite", overwrite=True)

You can also use your own image. Uncomment and run the following cell to use your own image.

1
# image = '/path/to/your/own/image.tif'

Display the downloaded image on the map.

1
2
3
m.layers[-1].visible = False
m.add_raster(image, layer_name="Image")
m

Initialize SAM class

Specify the file path to the model checkpoint. If it is not specified, the model will to downloaded to the working directory.

1
2
3
4
sam = SamGeo(
    model_type="vit_h",  # can be vit_h, vit_b, vit_l, vit_tiny
    sam_kwargs=None,
)

Automatic mask generation

Segment the image and save the results to a GeoTIFF file. Set unique=True to assign a unique ID to each object. 

1
sam.generate(image, output="masks.tif", foreground=True, unique=True)

1
sam.show_masks(cmap="binary_r")

Show the object annotations (objects with random color) on the map.

1
sam.show_anns(axis="off", alpha=1, output="annotations.tif")

Compare images with a slider.

1
2
3
4
5
6
leafmap.image_comparison(
    "satellite.tif",
    "annotations.tif",
    label1="Satellite Image",
    label2="Image Segmentation",
)

Add image to the map.

1
2
m.add_raster("annotations.tif", alpha=0.5, layer_name="Masks")
m

Convert the object annotations to vector format, such as GeoPackage, Shapefile, or GeoJSON.

1
sam.tiff_to_vector("masks.tif", "masks.gpkg")

Automatic mask generation options

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:

1
2
3
4
5
6
7
8
sam_kwargs = {
    "points_per_side": 32,
    "pred_iou_thresh": 0.86,
    "stability_score_thresh": 0.92,
    "crop_n_layers": 1,
    "crop_n_points_downscale_factor": 2,
    "min_mask_region_area": 100,
}

1
2
3
4
sam = SamGeo(
    model_type="vit_h",
    sam_kwargs=sam_kwargs,
)

1
sam.generate(image, output="masks2.tif", foreground=True)

1
sam.show_masks(cmap="binary_r")

1
sam.show_anns(axis="off", opacity=1, output="annotations2.tif")

Compare images with a slider.

1
2
3
4
5
6
leafmap.image_comparison(
    image,
    "annotations.tif",
    label1="Image",
    label2="Image Segmentation",
)

Overlay the annotations on the image and use the slider to change the opacity interactively.

1
overlay_images(image, "annotations2.tif", backend="TkAgg")