97 lines
4.0 KiB
Markdown
97 lines
4.0 KiB
Markdown
# water_V5.py：超大遥感影像水体分割（SAM3 + Rasterio 分块）
|
||
|
||
这个脚本用于在超大 GeoTIFF/遥感影像上分割水体（如 river/water body），重点在于在 16GB 显存条件下平衡速度与显存：
|
||
- 影像按“子区域”切分处理，避免一次性处理整图
|
||
- 每个子区域先做一次低分辨率粗分割定位水体，再只对“边缘带”进行高分辨率分块精修
|
||
- 精修的概率融合与大数组操作尽量放在 CPU（float16），显著降低显存占用
|
||
|
||
对应脚本：[water_V5.py](file:///e:/code/sam3-main/sam3-main/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](file:///e:/code/sam3-main/sam3-main/water_V5.py)，修改下面参数：
|
||
- `image_path`：输入 GeoTIFF
|
||
- `mask_output_path`：输出掩膜 GeoTIFF
|
||
- `prompt`：文本提示（例如 `"water body"` / `"river"` / `"water"`）
|
||
2. 运行：
|
||
|
||
```bash
|
||
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`
|
||
|