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 警告，不影响推理结果。

运行方式

1. 打开 water_V5.py，修改下面参数：
   • image_path：输入 GeoTIFF
   • mask_output_path：输出掩膜 GeoTIFF
   • prompt：文本提示（例如 "water body" / "river" / "water"）
2. 运行：

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