Segment Anything Model for Geospatial Data¶

This notebook shows how to use segment satellite imagery using the Segment Anything Model (SAM) with a few lines of code.

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.

Install dependencies¶

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

1

# %pip install leafmap

1

# %pip install segment-geospatial localtileserver

Import libraries¶

1
2
3

import os
import leafmap
from samgeo import SamGeo, tms_to_geotiff, get_basemaps

Create an interactive map¶

1
2
3

m = leafmap.Map(center=[29.676840, -95.369222], zoom=19)
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 = [-95.3704, 29.6762, -95.368, 29.6775]

Download map tiles¶

Download maps tiles and mosaic them into a single GeoTIFF file

1

image = "satellite.tif"

Besides the satellite basemap, you can use any of the following basemaps returned by the get_basemaps() function:

1

# get_basemaps().keys()

Specify the basemap as the source.

1

tms_to_geotiff(output=image, bbox=bbox, zoom=20, 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  # turn off the basemap
m.add_raster(image, layer_name="Image")
m

Initialize SAM class¶

1
2
3
4
5

sam = SamGeo(
    model_type="vit_h",
    checkpoint="sam_vit_h_4b8939.pth",
    sam_kwargs=None,
)

Segment the image¶

Set batch=True to segment the image in batches. This is useful for large images that cannot fit in memory.

1
2
3
4

mask = "segment.tif"
sam.generate(
    image, mask, batch=True, foreground=True, erosion_kernel=(3, 3), mask_multiplier=255
)

Polygonize the raster data¶

Save the segmentation results as a GeoPackage file.

1
2

vector = "segment.gpkg"
sam.tiff_to_gpkg(mask, vector, simplify_tolerance=None)

You can also save the segmentation results as any vector data format supported by GeoPandas.

1
2

shapefile = "segment.shp"
sam.tiff_to_vector(mask, shapefile)

Visualize the results¶

1
2
3
4
5
6
7
8

style = {
    "color": "#3388ff",
    "weight": 2,
    "fillColor": "#7c4185",
    "fillOpacity": 0.5,
}
m.add_vector(vector, layer_name="Vector", style=style)
m