# 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`