使用Segment-Geospatial项目实现HQ-SAM自动对象掩膜生成-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00162/article/details/148578650

使用Segment-Geospatial项目实现HQ-SAM自动对象掩膜生成

segment-geospatial A Python package for segmenting geospatial data with the Segment Anything Model (SAM) 项目地址: https://gitcode.com/gh_mirrors/se/segment-geospatial

概述

本文将介绍如何利用Segment-Geospatial项目中的High-Quality Segment Anything Model (HQ-SAM)技术，通过简单的代码实现图像中对象的自动分割。HQ-SAM是一种高质量的图像分割模型，能够准确识别和分割图像中的各种对象，特别适用于地理空间图像分析。

环境准备

在开始之前，需要确保已安装必要的依赖包。Segment-Geospatial项目提供了完整的工具链，可以方便地进行安装：

# 安装核心依赖包
%pip install segment-geospatial

安装完成后，导入必要的模块：

import os
import leafmap
from samgeo.hq_sam import (
    SamGeo,
    show_image,
    download_file,
    overlay_images,
    tms_to_geotiff,
)

创建交互式地图

使用leafmap库创建一个交互式地图，方便选择感兴趣区域：

# 初始化地图，设置中心点和缩放级别
m = leafmap.Map(center=[37.8713, -122.2580], zoom=17, height="800px")
m.add_basemap("SATELLITE")  # 添加卫星底图
m

在地图上通过绘制工具(多边形或矩形)选择感兴趣区域：

# 获取用户绘制的区域边界
if m.user_roi_bounds() is not None:
    bbox = m.user_roi_bounds()
else:
    # 默认区域(伯克利大学校园)
    bbox = [-122.2659, 37.8682, -122.2521, 37.8741]

获取图像数据

有两种方式获取待分析的图像数据：

下载示例卫星图像：

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

使用自定义图像(取消注释并修改路径)：

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

将图像添加到地图中显示：

m.layers[-1].visible = False  # 隐藏底图
m.add_raster(image, layer_name="Image")  # 添加图像图层
m

初始化HQ-SAM模型

HQ-SAM提供了多种模型类型，可根据需求选择：

sam = SamGeo(
    model_type="vit_h",  # 可选: vit_h(高精度), vit_b(平衡), vit_l(轻量), vit_tiny(超轻量)
    sam_kwargs=None,     # 自定义参数
)

自动掩膜生成

执行图像分割并保存结果：

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

参数说明：

foreground=True: 只保留前景对象
unique=True: 为每个对象分配唯一ID

可视化分割结果：

sam.show_masks(cmap="binary_r")  # 黑白显示掩膜

生成带颜色的对象标注：

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

结果对比与分析

使用滑块工具对比原始图像和分割结果：

leafmap.image_comparison(
    "satellite.tif",
    "annotations.tif",
    label1="Satellite Image",
    label2="Image Segmentation",
)

将分割结果叠加到地图上：

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

结果导出与应用

将分割结果转换为矢量格式(支持多种格式)：

sam.tiff_to_vector("masks.tif", "masks.gpkg")  # GeoPackage格式
# sam.tiff_to_vector("masks.tif", "masks.shp")  # Shapefile格式
# sam.tiff_to_vector("masks.tif", "masks.geojson")  # GeoJSON格式

高级参数调优

HQ-SAM提供了丰富的参数用于优化分割效果：

sam_kwargs = {
    "points_per_side": 32,  # 每边采样点数(增加可提高细节)
    "pred_iou_thresh": 0.86,  # IoU预测阈值
    "stability_score_thresh": 0.92,  # 稳定性分数阈值
    "crop_n_layers": 1,  # 裁剪层数
    "crop_n_points_downscale_factor": 2,  # 裁剪点降采样因子
    "min_mask_region_area": 100,  # 最小掩膜区域面积
}

# 使用优化参数重新初始化
sam = SamGeo(model_type="vit_h", sam_kwargs=sam_kwargs)

# 重新生成掩膜
sam.generate(image, output="masks2.tif", foreground=True)

实际应用建议

模型选择：对于高分辨率图像推荐使用vit_h，对实时性要求高的场景可选择vit_b或vit_l
参数调整：
- 增加points_per_side可提高小对象识别率
- 调整min_mask_region_area可过滤噪声
- 修改阈值参数可平衡精度和召回率
性能优化：
- 大图像可先裁剪后处理
- 批量处理时注意内存管理
- 考虑使用GPU加速