无人机河道航拍语义分割数据集 | 水利巡检、水体识别、洪涝监测、水资源AI分析数据集10330期

鸭蛋粉

205人浏览 · 2026-06-12 05:30:00

鸭蛋粉 · 2026-06-12 05:30:00 发布

无人机河道航拍语义分割数据集 | 水利巡检、水体识别、洪涝监测、水资源AI分析数据集10330期

摘要：一套面向无人机航拍河道场景的水体语义分割开源数据集，适配主流深度学习语义分割框架，可直接用于水体区域提取、常态化河道巡检、洪涝灾害范围评估、水资源动态监测等水利AI业务，标注规范、开箱即用，适配算法原型开发、模型迭代与工程落地。

标签：#无人机河道巡检 #水体语义分割 #水利AI #洪涝灾害监测 #水资源智能分析 #航拍图像分割 #深度学习数据集 #环境监测CV #河流区域识别

在这里插入图片描述

传统河道、流域、水资源监测长期依靠人工徒步、船只巡航完成巡查与水体范围统计，河道跨度大、沿岸地形复杂，部分浅滩、急流、偏远水域人员与船只难以抵达，巡检覆盖不全、作业风险高、人力与时间成本居高不下。在汛期洪涝灾害场景中，人工排查受淹区域速度慢、范围统计误差大，无法为应急指挥提供实时、精准的可视化数据支撑；同时，传统监测手段难以实现全天候、高频次流域动态追踪，无法满足现代化智慧水利、河湖长制的数字化管理需求。

如今无人机航拍+语义分割已成为水利行业智能化升级的主流路线，依靠AI自动分割水体区域，可快速划定河流、湖泊、积水、受淹区域边界。但目前行业普遍存在专业水体分割标注数据稀缺、通用分割数据集与河道航拍场景不匹配、标注格式混乱、复杂水面场景样本不足等问题，大幅抬高了水利AI算法落地门槛。

本次分享的河道航拍语义分割数据集，聚焦真实野外河道、自然水域场景，专门针对水体区域做精细化像素级标注，格式兼容主流语义分割框架，数据场景丰富，能够有效解决水利视觉算法研发过程中数据短缺、场景适配差的痛点。

在这里插入图片描述

二、数据集整体详情

2.1 核心基础信息

本数据集为纯语义分割任务专用标注数据集，区别于图像分类、目标检测数据集，以像素级水体区域分割为核心标注目标，所有数据均来源于无人机实地航拍，贴合真实水利巡检作业画面。

标注目标：单类别语义分割，仅区分水体区域与背景区域，标注聚焦水域边界、水面范围，精准适配河道识别、积水提取、洪区划定等核心业务。
数据规模：总计644张高清航拍标注图像，数据体量适中，兼顾模型训练效率与泛化能力，非常适合语义分割算法入门验证、中小型模型训练、企业原型项目开发、高校水利AI实训。
场景覆盖：样本涵盖不同光照、不同岸线形态、不同水流状态的河道场景，有效降低单一环境数据造成的模型过拟合问题，提升模型在野外复杂水域的鲁棒性。
数据格式：采用计算机视觉语义分割任务通用标准格式，图像与掩码标签一一对应，兼容U-Net、PSPNet、SegFormer、YOLO-Seg等主流深度学习框架，无需复杂格式转换，下载后可直接加载训练。

2.2 核心落地应用价值

算法研发支撑：为水体语义分割模型提供标准化训练、验证数据，助力开发者调试网络结构、优化损失函数、对比不同分割算法性能。
常态化水资源监测：结合无人机航拍，自动分割、统计流域水体面积，实现水资源分布常态化盘点与动态追踪，替代人工丈量统计。
洪涝灾害应急处置：汛期快速分割识别积水、淹没区域范围，输出精准水域边界数据，辅助灾情评估、人员转移、抢险路线规划。
生态环境监测：应用于河湖生态管护、水域岸线监管、黑臭水体范围圈定等场景，推动计算机视觉技术在水利、环保垂直领域规模化落地。

2.3 数据集核心优势

场景垂直化：原生无人机河道航拍数据，水面反光、岸线遮挡、远景水域等真实干扰场景全覆盖，远优于通用户外分割数据集；
标注轻量化：单类别分割任务，学习门槛低、训练收敛速度快，新手可快速上手；
算力友好：644张图像体量小，普通家用GPU、边缘算力设备即可完成全流程训练与推理，无需高性能集群；
全框架兼容：标准掩码标注格式，适配PyTorch、TensorFlow两大主流生态，适配绝大多数开源语义分割项目。

三、深度学习实战代码（附河道场景专属注释）

基于水利行业语义分割常用框架 U-Net 与 Ultralytics YOLO-Seg 编写全套代码，包含环境部署、数据集加载、模型训练、图像推理、结果可视化全流程，所有代码结合河道航拍、水面分割场景做针对性注释与参数调优。

3.1 运行环境依赖安装

适配 Python 3.8 ~ 3.11，支持 Windows / Linux 系统，CUDA 10.2 及以上版本均可运行。

# 安装语义分割核心依赖库
pip install torch torchvision  # 深度学习基础框架
pip install segmentation-models-pytorch==0.3.3  # 快速调用U-Net等经典分割模型
pip install opencv-python pillow numpy  # 图像处理、矩阵运算
pip install albumentations  # 图像增强，适配水面反光、光照变化场景
pip install tqdm  # 训练进度可视化

3.2 数据集标准目录结构（通用语义分割格式）

遵循行业通用目录规范，图像与掩码标签文件名一一对应，解压后直接使用，无需重构路径：

water_segment_dataset/
├── images/                # 无人机航拍原图（总计644张）
│   ├── train/             # 训练集图像
│   ├── val/               # 验证集图像
│   └── test/              # 测试集图像
├── masks/                 # 语义分割掩码标签（单通道PNG，像素值代表类别）
│   ├── train/
│   ├── val/
│   └── test/
└── water_config.yaml      # 数据集配置文件（类别、路径、超参数）

3.3 数据集配置文件 `water_config.yaml`

针对水体单类别分割场景编写配置文件，适配YOLO-Seg系列框架：

# 无人机河道水体语义分割 - 数据集配置文件
# 数据集根路径（根据本地实际存放路径修改）
path: ./water_segment_dataset
# 训练、验证、测试集图像路径
train: images/train
val: images/val
test: images/test

# 分割类别数量：2类（背景、水体）
nc: 2
# 类别名称映射
names:
  0: 背景
  1: 水体

# 河道航拍专属数据增强（针对水面反光、逆光、画面翻转优化）
fliplr: 0.5    # 左右翻转，提升模型对不同走向河道的适应性
hsv_h: 0.015   # 色相微调，适配早晚光照变化
hsv_s: 0.7     # 饱和度微调，削弱水面反光干扰

3.4 自定义数据集加载脚本 `water_dataset.py`

基于PyTorch Dataset 封装数据集类，适配U-Net训练，适配航拍图像与掩码读取逻辑：

"""
河道水体分割 - 自定义数据集加载类
适配场景：无人机航拍河道图像、像素级水体掩码标签
适配框架：PyTorch + segmentation-models-pytorch
"""
import os
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
import albumentations as A
from albumentations.pytorch import ToTensorV2

class WaterSegmentDataset(Dataset):
    def __init__(self, image_dir, mask_dir, transform=None):
        # 图像文件夹、掩码标签文件夹路径
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.transform = transform
        # 读取所有图像文件名，保证图像与标签顺序一致
        self.images = sorted([f for f in os.listdir(image_dir) if f.endswith(('.jpg', '.png'))])

    def __len__(self):
        # 返回数据集总样本数
        return len(self.images)

    def __getitem__(self, idx):
        # 拼接单张图像与标签完整路径
        img_name = self.images[idx]
        img_path = os.path.join(self.image_dir, img_name)
        mask_path = os.path.join(self.mask_dir, img_name.replace(".jpg", ".png"))

        # 读取航拍图像与分割掩码
        image = np.array(Image.open(img_path).convert("RGB"))
        mask = np.array(Image.open(mask_path), dtype=np.float32)

        # 掩码二值化：1=水体，0=背景（适配单类别水体分割）
        mask[mask > 0] = 1.0

        # 执行数据增强（针对水面反光、光照干扰做增强）
        if self.transform:
            aug = self.transform(image=image, mask=mask)
            image = aug["image"]
            mask = aug["mask"]

        return image, mask

# 定义训练集数据增强策略（河道场景专属）
def get_train_aug():
    return A.Compose([
        A.Resize(640, 640),  # 统一图像尺寸，适配航拍宽画幅
        A.HorizontalFlip(p=0.5),  # 随机左右翻转
        A.RandomBrightnessContrast(p=0.3),  # 明暗、对比度调整，削弱水面反光
        A.Normalize(),
        ToTensorV2()
    ])

# 定义验证集数据增强（仅归一化，不做随机变换）
def get_val_aug():
    return A.Compose([
        A.Resize(640, 640),
        A.Normalize(),
        ToTensorV2()
    ])

3.5 U-Net 模型训练脚本 `train_unet_water.py`

选用水利分割场景常用的U-Net模型，结合644张样本特性与河道航拍场景调优超参，防止小数据集过拟合：

"""
无人机河道水体分割 - U-Net训练脚本
数据集规模：644张航拍图像
适配场景：水面反光、岸线遮挡、野外河道全域水体分割
核心优化：小数据集防过拟合、水面特征强化
"""
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import segmentation_models_pytorch as smp
from tqdm import tqdm
from water_dataset import WaterSegmentDataset, get_train_aug, get_val_aug

# ===================== 全局配置（河道场景专属参数）=====================
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 数据集路径
TRAIN_IMG = "./water_segment_dataset/images/train"
TRAIN_MASK = "./water_segment_dataset/masks/train"
VAL_IMG = "./water_segment_dataset/images/val"
VAL_MASK = "./water_segment_dataset/masks/val"
# 训练超参（644张小数据集，降低batch、减少轮次防过拟合）
BATCH_SIZE = 8
EPOCHS = 50
LEARNING_RATE = 1e-4
# =====================================================================

def train():
    # 1. 加载训练集、验证集
    train_dataset = WaterSegmentDataset(TRAIN_IMG, TRAIN_MASK, transform=get_train_aug())
    val_dataset = WaterSegmentDataset(VAL_IMG, VAL_MASK, transform=get_val_aug())

    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

    # 2. 初始化U-Net模型，骨干网络选用efficientnet-b0（轻量化，适配边缘部署）
    model = smp.Unet(
        encoder_name="efficientnet-b0",
        encoder_weights="imagenet",  # 使用预训练权重，小数据集加速收敛
        in_channels=3,
        classes=1,  # 输出单通道掩码，二分类分割
        activation="sigmoid"
    ).to(DEVICE)

    # 3. 损失函数 + 优化器（水体分割二分类任务首选BCE损失）
    loss_fn = nn.BCELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

    # 4. 开始迭代训练
    for epoch in range(EPOCHS):
        # 训练阶段
        model.train()
        train_loss = 0.0
        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{EPOCHS} Train")
        for imgs, masks in pbar:
            imgs, masks = imgs.to(DEVICE), masks.to(DEVICE)
            optimizer.zero_grad()
            outputs = model(imgs)
            loss = loss_fn(outputs, masks.unsqueeze(1))
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
            pbar.set_postfix({"loss": loss.item()})

        # 验证阶段
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            pbar_val = tqdm(val_loader, desc=f"Epoch {epoch+1}/{EPOCHS} Val")
            for imgs, masks in pbar_val:
                imgs, masks = imgs.to(DEVICE), masks.to(DEVICE)
                outputs = model(imgs)
                loss = loss_fn(outputs, masks.unsqueeze(1))
                val_loss += loss.item()

        # 打印每轮平均损失
        avg_train_loss = train_loss / len(train_loader)
        avg_val_loss = val_loss / len(val_loader)
        print(f"平均训练损失: {avg_train_loss:.4f} | 平均验证损失: {avg_val_loss:.4f}")

    # 保存训练完成的模型权重
    torch.save(model.state_dict(), "./water_unet_best.pth")
    print("模型训练完成，权重已保存为 water_unet_best.pth")

if __name__ == "__main__":
    train()

3.6 图像推理与可视化脚本 `predict_water.py`

加载训练好的模型，对无人机航拍图像做水体分割推理，输出水体掩码与叠加效果图，适配巡检业务可视化需求：

"""
河道水体分割推理脚本
功能：单张/批量航拍图像水体提取、分割结果可视化
落地场景：无人机实时巡检画面分析、洪涝区域离线统计
"""
import torch
import cv2
import numpy as np
import os
import segmentation_models_pytorch as smb
from PIL import Image

# 全局配置
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_PATH = "./water_unet_best.pth"  # 训练好的权重路径
TEST_IMG_DIR = "./water_segment_dataset/images/test"  # 测试图像目录
SAVE_RESULT_DIR = "./segment_result"  # 结果保存目录
os.makedirs(SAVE_RESULT_DIR, exist_ok=True)

# 加载模型
model = smb.Unet(
    encoder_name="efficientnet-b0",
    encoder_weights=None,
    in_channels=3,
    classes=1,
    activation="sigmoid"
).to(DEVICE)
model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
model.eval()

def predict_single_image(img_path):
    # 读取航拍图像
    image = cv2.imread(img_path)
    h, w = image.shape[:2]
    # 图像预处理
    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    img_resize = cv2.resize(img_rgb, (640, 640))
    img_tensor = torch.from_numpy(img_resize).permute(2, 0, 1).float() / 255.0
    img_tensor = img_tensor.unsqueeze(0).to(DEVICE)

    # 模型推理（阈值0.5区分水体与背景，适配水面模糊边界）
    with torch.no_grad():
        output = model(img_tensor)
        pred_mask = (output > 0.5).float().squeeze().cpu().numpy()

    # 掩码还原原图尺寸
    pred_mask = cv2.resize(pred_mask, (w, h))
    # 生成蓝色掩码图层，叠加到原图（直观展示水体区域）
    mask_color = np.zeros_like(image)
    mask_color[pred_mask == 1] = [255, 0, 0]  # 蓝色标注水体
    # 原图+掩码融合
    result = cv2.addWeighted(image, 0.7, mask_color, 0.3, 0)
    return result

# 批量处理测试图像
if __name__ == "__main__":
    for img_name in os.listdir(TEST_IMG_DIR):
        if img_name.endswith((".jpg", ".png")):
            img_full_path = os.path.join(TEST_IMG_DIR, img_name)
            seg_result = predict_single_image(img_full_path)
            save_path = os.path.join(SAVE_RESULT_DIR, img_name)
            cv2.imwrite(save_path, seg_result)
    print(f"批量水体分割完成，结果已保存至 {SAVE_RESULT_DIR}")