Segment Anything模型用于地理空间数据

news2024/10/6 22:21:27

原文地址：https://samgeo.gishub.org/examples/satellite/

此笔记本通过几行代码展示了如何使用 Segment Anything Model (SAM) 来使用分段卫星图像。确保为此jupyter notebook使用 GPU 运行时。

安装依赖

取消注释并运行以下单元格以安装所需的依赖项。

# %pip install segment-geospatial leafmap localtileserver

导入库

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

创建地图应用

m = leafmap.Map(center=[29.676840, -95.369222], zoom=19)
m.add_basemap('SATELLITE')
m
if m.user_roi_bounds() is not None:
    bbox = m.user_roi_bounds()
else:
    bbox = [-95.3704, 29.6762, -95.368, 29.6775]

下载地图图块

下载地图图块并将它们镶嵌到单个 GeoTIFF 文件中

image = 'satellite.tif'
# image = '/path/to/your/own/image.tif'
 tms_to_geotiff(output=image, bbox=bbox, zoom=20, source='Satellite')
 m.add_raster(image, layer_name='Image')

初始化 SAM 类

out_dir = os.path.join(os.path.expanduser('~'), 'Downloads')
checkpoint = os.path.join(out_dir, 'sam_vit_h_4b8939.pth')
device = 'cuda' if torch.cuda.is_available() else 'cpu'
sam = SamGeo(
    checkpoint=checkpoint,
    model_type='vit_h',
    device=device,
    erosion_kernel=(3, 3),
    mask_multiplier=255,
    sam_kwargs=None,
)

分割图像

mask = 'segment.tiff'
sam.generate(image, mask)

多边形化栅格数据

#将分割结果保存为 GeoPackage 文件。
vector = 'segment.gpkg'
sam.tiff_to_gpkg(mask, vector, simplify_tolerance=None)
#您还可以将分割结果保存为 GeoPandas 支持的任何矢量数据格式。
shapefile = 'segment.shp'
sam.tiff_to_vector(mask, shapefile)

可视化结果

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