基于深度学习的水体提取

背景

本科做毕设时，我用水体指数（NDWI）来提取水体边界。这个方法简单，但一旦遇到地形复杂、地物干扰多的场景，效果就不太理想。这段时间在学习深度学习，就想顺手试一下。我选了 DeepWaterMap 来练手。它是针对多光谱影像做地表水体分割的网络，流程不复杂，比较适合快速上手。

1. 模型和输入数据

DeepWaterMap 期望输入为多光谱影像，包含以下波段：

B2: Blue
B3: Green
B4: Red
B5: Near Infrared (NIR)
B6: Shortwave Infrared 1 (SWIR1)
B7: Shortwave Infrared 2 (SWIR2)

模型结构比较简洁：

网络主要由三个模块组成：

• 下采样单元：通过步幅卷积做空间降采样。
• 瓶颈单元：输出与输入维度一致的特征图。
• 上采样单元：先做 depth-to-space，再通过卷积恢复空间分辨率。

整体看下来，这个结构在计算量和分割效果之间做了一个比较好的平衡。

2. 环境配置和推理

创建新的 conda 环境后安装依赖：

pip install tensorflow==1.12.0 tifffile opencv-python

下载好权重文件后，就可以直接跑推理：

python inference.py --checkpoint_path checkpoints/cp.135.ckpt \
    --image_path sample_data/sentinel2_example.tif --save_path water_map.png

测试我用的是 Sentinel-2 影像。原始数据需要先做波段组合，按模型要求把波段顺序整理好，再作为输入去推理。

3. 批量推理

我把原始脚本稍微改了一下，用来批量推理：

▶

批量推理示例脚本

import os
import argparse
import deepwatermap
import tifffile as tiff
import numpy as np
import cv2
import matplotlib.pyplot as plt
from download_data import download_from_gdrive
from projection import set_projection
import glob
import tensorflow as tf

def find_padding(v, divisor=32):
    v_divisible = max(divisor, int(divisor * np.ceil( v / divisor )))
    total_pad = v_divisible - v
    pad_1 = total_pad // 2
    pad_2 = total_pad - pad_1
    return pad_1, pad_2

def predict(checkpoint_path, image_path, save_path):
    image = tiff.imread(image_path)
    pad_r = find_padding(image.shape[0])
    pad_c = find_padding(image.shape[1])
    image = np.pad(image, ((pad_r[0], pad_r[1]), (pad_c[0], pad_c[1]), (0, 0)), 'reflect')
    image = image.astype(np.float32)
    image = image - np.min(image)
    image = image / np.maximum(np.max(image), 1)
    image = np.expand_dims(image, axis=0)
    dwm = model.predict(image)
    dwm = np.squeeze(dwm)
    dwm = dwm[pad_r[0]:-pad_r[1], pad_c[0]:-pad_c[1]]
    dwm = 1./(1+np.exp(-(16*(dwm-0.5))))
    dwm = np.clip(dwm, 0, 1)
    cv2.imwrite(save_path, dwm * 255)

def load_model_gpus(checkpoint_path):
    strategy = tf.distribute.MirroredStrategy()
    with strategy.scope():
        model = deepwatermap.model()
        model.load_weights(checkpoint_path) 
    return model

if __name__ == '__main__':
    work_dir = os.path.dirname(os.path.abspath(__file__))
    in_image_folder = "sample_data"
    use_png_ext = True
    output_dir = os.path.join(work_dir, "results")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    image_paths = glob.glob(os.path.join(work_dir, in_image_folder, '*.tif'))

    # 使用GPU进行计算
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            # 设置GPU内存自增长
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
        except RuntimeError as e:
            print(e)

    batch_size = 4  # 每个GPU一次读取的影像数量
    checkpoint_path = os.path.join(work_dir, "checkpoints/cp.135.ckpt")
    model = load_model_gpus(checkpoint_path)
    num_gpus = len(gpus)

    # 平均分配影像给每个GPU
    image_per_gpu = len(image_paths) // num_gpus
    for i in range(num_gpus):
        gpu_image_paths = image_paths[i * image_per_gpu : (i+1) * image_per_gpu]
        # 如果最后一个GPU处理的影像数与其他GPU不同，则把剩余的影像全交给最后一个GPU处理
        if i == num_gpus - 1 and len(image_paths) % num_gpus != 0:
            gpu_image_paths += image_paths[num_gpus * image_per_gpu:]
        for in_image_path in gpu_image_paths:
            basename, extension = os.path.splitext(os.path.basename(in_image_path))
            if use_png_ext:
                extension = ".png"
            out_image_name = basename + "_water" + extension
            out_image_path = os.path.join(output_dir, out_image_name)
            predict(checkpoint_path, in_image_path, out_image_path)
            if not use_png_ext:
                set_projection(out_image_path, template_raster=in_image_path)
            print("\nOutput path: {}".format(out_image_path))

4. 结果和感受

输出结果可以通过阈值进一步调节，设置为 0.5 效果还行，0.3 会更宽松一些，0.7 则更严格一些。

从实际效果看，模型对主要水体区域识别得比较稳，边界也更清楚一些。实际上现在很多的地物分割只用 RGB，结果也不差；如果用多光谱输入，光谱信息更完整，分割质量还能再往上提一点。

5. 小结

• 传统指数法适合快速出图，但在复杂场景下鲁棒性会差一些。
• 深度学习方法在边界表达和复杂背景区分上，确实更有潜力。
• 数据预处理（波段顺序、归一化、阈值设置）对最终结果影响很大。