4.0 KiB
4.0 KiB
water_V5.py：超大遥感影像水体分割（SAM3 + Rasterio 分块）
这个脚本用于在超大 GeoTIFF/遥感影像上分割水体（如 river/water body），重点在于在 16GB 显存条件下平衡速度与显存：
- 影像按“子区域”切分处理，避免一次性处理整图
- 每个子区域先做一次低分辨率粗分割定位水体，再只对“边缘带”进行高分辨率分块精修
- 精修的概率融合与大数组操作尽量放在 CPU（float16），显著降低显存占用
对应脚本：water_V5.py
功能概览
- 读取 GeoTIFF 子区域（Rasterio Window/MemoryFile）
- overview 粗分割（低分辨率一次推理）
- 由粗分割生成边缘带（ring/band）
- 仅对边缘带覆盖到的 tile 做精修推理（重叠窗口融合）
- 后处理：内部 NoData 空洞填充、连通域面积过滤
- 输出单通道 uint8 掩膜 GeoTIFF（像元值 0/1）
环境依赖
- Python + PyTorch（GPU 可选但推荐）
- rasterio
- numpy
- PIL (Pillow)
- scipy
- tqdm
提示：运行时可能看到 pkg_resources is deprecated 警告，不影响推理结果。
运行方式
- 打开 water_V5.py，修改下面参数：
image_path：输入 GeoTIFFmask_output_path：输出掩膜 GeoTIFFprompt：文本提示（例如"water body"/"river"/"water"）
- 运行：
python e:\code\sam3-main\sam3-main\water_V5.py
输出 mask_output_path 为 uint8 单通道，像元值为 0/1，并保留输入影像的空间参考（transform/CRS）。
关键参数说明（建议优先调这几个）
在脚本“参数设置”区域：
coarse_read_max_side：overview 粗分割输入的最大边长（越大越准/越慢，越小越快/越粗）- 建议：
1024 ~ 1536（16GB 显存一般可从 1200 开始）
- 建议：
tile_size：精修分块大小（原图像素）- 建议：
1536（更稳）或2048（更快但更吃内存/更容易慢）
- 建议：
overlap：精修分块重叠（减少拼接缝）- 建议：
128（更快）或256（边缘更稳但慢）
- 建议：
band_radius：边缘带宽度（越大精修范围越大，速度越慢）- 建议：
48 ~ 96
- 建议：
coarse_threshold/final_threshold：粗分割/最终阈值- 通常
0.5可用；若漏检多，可适当降低阈值（例如0.4）
- 通常
分区域参数：
num_splits_y、num_splits_x：将整图拆成多少个子区域处理region_overlap：子区域之间重叠，避免区域边界出现断裂
后处理参数：
min_area：移除小碎片的最小连通域面积（像素）keep_largest_only：只保留最大连通域（适合只想要“最大水体”的场景）
性能与显存的取舍建议（16GB）
优先级从高到低：
- 优先调小
tile_size（例如 2048 → 1536），通常能显著降低显存与峰值计算量 overlap从 256 降到 128，可明显加速（代价是拼接缝更可能出现）band_radius适当减小，让精修只发生在更窄的边缘范围coarse_read_max_side适度提高可减少误分/漏分，但会更慢
脚本已做的关键优化：
- 精修 tile 的“是否需要处理”的判断在 CPU 上完成，避免每个 tile 强制 GPU 同步
- per-tile 的强度拉伸（percentile）改为“子区域级别一次计算并复用”
- 推理使用 autocast（float16），通常更省显存
输出说明
- 输出掩膜为
uint8单通道：- 0：非水体
- 1：水体
- 输出 GeoTIFF 继承输入的 profile，并设置
compress='lzw'
常见问题
-
运行报
AssertionError（RoPE 相关）- 该模型对输入分辨率有约束。脚本默认
coarse_resolution = 1008、fine_resolution = 1008，不要随意改成 768/512 等。
- 该模型对输入分辨率有约束。脚本默认
-
输出全黑/全白
- 先换
prompt（如"water"/"river"），再调coarse_threshold/final_threshold - 若影像波段不是 RGB，脚本会取前三个波段做可视输入；需要更合适的波段组合时应自行调整
_bands_to_pil
- 先换