无人机河道航拍语义分割数据集 | 水利巡检、水体识别、洪涝监测、水资源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 及以上版本均可运行。# 安装语义分割核心依赖库pipinstalltorch torchvision# 深度学习基础框架pipinstallsegmentation-models-pytorch0.3.3# 快速调用U-Net等经典分割模型pipinstallopencv-python pillow numpy# 图像处理、矩阵运算pipinstallalbumentations# 图像增强适配水面反光、光照变化场景pipinstalltqdm# 训练进度可视化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/trainval:images/valtest: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 importosimportnumpyasnpfromPILimportImagefromtorch.utils.dataimportDatasetimportalbumentationsasAfromalbumentations.pytorchimportToTensorV2classWaterSegmentDataset(Dataset):def__init__(self,image_dir,mask_dir,transformNone):# 图像文件夹、掩码标签文件夹路径self.image_dirimage_dir self.mask_dirmask_dir self.transformtransform# 读取所有图像文件名保证图像与标签顺序一致self.imagessorted([fforfinos.listdir(image_dir)iff.endswith((.jpg,.png))])def__len__(self):# 返回数据集总样本数returnlen(self.images)def__getitem__(self,idx):# 拼接单张图像与标签完整路径img_nameself.images[idx]img_pathos.path.join(self.image_dir,img_name)mask_pathos.path.join(self.mask_dir,img_name.replace(.jpg,.png))# 读取航拍图像与分割掩码imagenp.array(Image.open(img_path).convert(RGB))masknp.array(Image.open(mask_path),dtypenp.float32)# 掩码二值化1水体0背景适配单类别水体分割mask[mask0]1.0# 执行数据增强针对水面反光、光照干扰做增强ifself.transform:augself.transform(imageimage,maskmask)imageaug[image]maskaug[mask]returnimage,mask# 定义训练集数据增强策略河道场景专属defget_train_aug():returnA.Compose([A.Resize(640,640),# 统一图像尺寸适配航拍宽画幅A.HorizontalFlip(p0.5),# 随机左右翻转A.RandomBrightnessContrast(p0.3),# 明暗、对比度调整削弱水面反光A.Normalize(),ToTensorV2()])# 定义验证集数据增强仅归一化不做随机变换defget_val_aug():returnA.Compose([A.Resize(640,640),A.Normalize(),ToTensorV2()])3.5 U-Net 模型训练脚本train_unet_water.py选用水利分割场景常用的U-Net模型结合644张样本特性与河道航拍场景调优超参防止小数据集过拟合 无人机河道水体分割 - U-Net训练脚本 数据集规模644张航拍图像 适配场景水面反光、岸线遮挡、野外河道全域水体分割 核心优化小数据集防过拟合、水面特征强化 importtorchimporttorch.nnasnnfromtorch.utils.dataimportDataLoaderimportsegmentation_models_pytorchassmpfromtqdmimporttqdmfromwater_datasetimportWaterSegmentDataset,get_train_aug,get_val_aug# 全局配置河道场景专属参数DEVICEtorch.device(cudaiftorch.cuda.is_available()elsecpu)# 数据集路径TRAIN_IMG./water_segment_dataset/images/trainTRAIN_MASK./water_segment_dataset/masks/trainVAL_IMG./water_segment_dataset/images/valVAL_MASK./water_segment_dataset/masks/val# 训练超参644张小数据集降低batch、减少轮次防过拟合BATCH_SIZE8EPOCHS50LEARNING_RATE1e-4# deftrain():# 1. 加载训练集、验证集train_datasetWaterSegmentDataset(TRAIN_IMG,TRAIN_MASK,transformget_train_aug())val_datasetWaterSegmentDataset(VAL_IMG,VAL_MASK,transformget_val_aug())train_loaderDataLoader(train_dataset,batch_sizeBATCH_SIZE,shuffleTrue,num_workers2)val_loaderDataLoader(val_dataset,batch_sizeBATCH_SIZE,shuffleFalse,num_workers2)# 2. 初始化U-Net模型骨干网络选用efficientnet-b0轻量化适配边缘部署modelsmp.Unet(encoder_nameefficientnet-b0,encoder_weightsimagenet,# 使用预训练权重小数据集加速收敛in_channels3,classes1,# 输出单通道掩码二分类分割activationsigmoid).to(DEVICE)# 3. 损失函数 优化器水体分割二分类任务首选BCE损失loss_fnnn.BCELoss()optimizertorch.optim.Adam(model.parameters(),lrLEARNING_RATE)# 4. 开始迭代训练forepochinrange(EPOCHS):# 训练阶段model.train()train_loss0.0pbartqdm(train_loader,descfEpoch{epoch1}/{EPOCHS}Train)forimgs,masksinpbar:imgs,masksimgs.to(DEVICE),masks.to(DEVICE)optimizer.zero_grad()outputsmodel(imgs)lossloss_fn(outputs,masks.unsqueeze(1))loss.backward()optimizer.step()train_lossloss.item()pbar.set_postfix({loss:loss.item()})# 验证阶段model.eval()val_loss0.0withtorch.no_grad():pbar_valtqdm(val_loader,descfEpoch{epoch1}/{EPOCHS}Val)forimgs,masksinpbar_val:imgs,masksimgs.to(DEVICE),masks.to(DEVICE)outputsmodel(imgs)lossloss_fn(outputs,masks.unsqueeze(1))val_lossloss.item()# 打印每轮平均损失avg_train_losstrain_loss/len(train_loader)avg_val_lossval_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加载训练好的模型对无人机航拍图像做水体分割推理输出水体掩码与叠加效果图适配巡检业务可视化需求 河道水体分割推理脚本 功能单张/批量航拍图像水体提取、分割结果可视化 落地场景无人机实时巡检画面分析、洪涝区域离线统计 importtorchimportcv2importnumpyasnpimportosimportsegmentation_models_pytorchassmbfromPILimportImage# 全局配置DEVICEtorch.device(cudaiftorch.cuda.is_available()elsecpu)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_okTrue)# 加载模型modelsmb.Unet(encoder_nameefficientnet-b0,encoder_weightsNone,in_channels3,classes1,activationsigmoid).to(DEVICE)model.load_state_dict(torch.load(MODEL_PATH,map_locationDEVICE))model.eval()defpredict_single_image(img_path):# 读取航拍图像imagecv2.imread(img_path)h,wimage.shape[:2]# 图像预处理img_rgbcv2.cvtColor(image,cv2.COLOR_BGR2RGB)img_resizecv2.resize(img_rgb,(640,640))img_tensortorch.from_numpy(img_resize).permute(2,0,1).float()/255.0img_tensorimg_tensor.unsqueeze(0).to(DEVICE)# 模型推理阈值0.5区分水体与背景适配水面模糊边界withtorch.no_grad():outputmodel(img_tensor)pred_mask(output0.5).float().squeeze().cpu().numpy()# 掩码还原原图尺寸pred_maskcv2.resize(pred_mask,(w,h))# 生成蓝色掩码图层叠加到原图直观展示水体区域mask_colornp.zeros_like(image)mask_color[pred_mask1][255,0,0]# 蓝色标注水体# 原图掩码融合resultcv2.addWeighted(image,0.7,mask_color,0.3,0)returnresult# 批量处理测试图像if__name____main__:forimg_nameinos.listdir(TEST_IMG_DIR):ifimg_name.endswith((.jpg,.png)):img_full_pathos.path.join(TEST_IMG_DIR,img_name)seg_resultpredict_single_image(img_full_path)save_pathos.path.join(SAVE_RESULT_DIR,img_name)cv2.imwrite(save_path,seg_result)print(f批量水体分割完成结果已保存至{SAVE_RESULT_DIR})四、数据集使用策略 场景化深度思考4.1 数据划分建议数据集共644张航拍图像结合语义分割任务特性与小样本训练规则推荐划分比例训练集70%451张用于模型主体训练验证集15%97张用于训练过程中监控精度、防止过拟合测试集15%96张用于模型最终效果评估与业务仿真测试。划分时保证不同光照、不同河道形态样本均匀分布避免局部场景数据集中。4.2 河道航拍分割场景调优经验应对水面反光干扰水面强反光是河道分割最大难点训练阶段开启亮度、对比度随机变换推理时可搭配图像滤波预处理能明显提升边界分割精度。小样本防过拟合本数据集仅644张图像训练时建议选用轻量化骨干网络、降低BatchSize、配合早停机制不要使用过深的网络结构。边缘部署适配U-Net轻量化版本推理速度快训练后的模型可直接部署在无人机机载端、水利边缘网关实现航拍实时分割端到端巡检。边界优化河道水岸交界区域像素模糊可在后处理中使用形态学运算优化分割边缘提升水域面积统计精度。4.3 技术拓展方向多任务融合在水体分割基础上叠加目标检测同时实现水域分割河道漂浮物、违建船只识别打造一体化智能巡检系统。时序分析结合多时段航拍时序图像分析水体面积变化实现流域水位、水量动态监测。跨场景迁移基于本数据集做预训练再用本地流域数据微调快速适配不同地域河道、湖泊、水库场景。五、总结本套无人机河道航拍水体语义分割数据集聚焦水利行业刚需场景644张标注样本场景丰富、标注规范、格式通用完美适配语义分割算法训练与水利AI工程落地。配套U-Net全流程代码覆盖数据加载、训练、推理、可视化从算法研发到业务部署形成完整链路。该数据集既适合AI学习者入门语义分割任务也可作为水利、环保企业智慧巡检、洪涝监测项目的基础数据底座能够大幅降低水利计算机视觉项目的研发成本与周期。附加标签#U-Net实战 #语义分割开源 #智慧水利 #航拍图像AI分析 #洪涝区域识别
无人机河道航拍语义分割数据集 | 水利巡检、水体识别、洪涝监测、水资源AI分析数据集10330期
无人机河道航拍语义分割数据集 | 水利巡检、水体识别、洪涝监测、水资源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 及以上版本均可运行。# 安装语义分割核心依赖库pipinstalltorch torchvision# 深度学习基础框架pipinstallsegmentation-models-pytorch0.3.3# 快速调用U-Net等经典分割模型pipinstallopencv-python pillow numpy# 图像处理、矩阵运算pipinstallalbumentations# 图像增强适配水面反光、光照变化场景pipinstalltqdm# 训练进度可视化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/trainval:images/valtest: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 importosimportnumpyasnpfromPILimportImagefromtorch.utils.dataimportDatasetimportalbumentationsasAfromalbumentations.pytorchimportToTensorV2classWaterSegmentDataset(Dataset):def__init__(self,image_dir,mask_dir,transformNone):# 图像文件夹、掩码标签文件夹路径self.image_dirimage_dir self.mask_dirmask_dir self.transformtransform# 读取所有图像文件名保证图像与标签顺序一致self.imagessorted([fforfinos.listdir(image_dir)iff.endswith((.jpg,.png))])def__len__(self):# 返回数据集总样本数returnlen(self.images)def__getitem__(self,idx):# 拼接单张图像与标签完整路径img_nameself.images[idx]img_pathos.path.join(self.image_dir,img_name)mask_pathos.path.join(self.mask_dir,img_name.replace(.jpg,.png))# 读取航拍图像与分割掩码imagenp.array(Image.open(img_path).convert(RGB))masknp.array(Image.open(mask_path),dtypenp.float32)# 掩码二值化1水体0背景适配单类别水体分割mask[mask0]1.0# 执行数据增强针对水面反光、光照干扰做增强ifself.transform:augself.transform(imageimage,maskmask)imageaug[image]maskaug[mask]returnimage,mask# 定义训练集数据增强策略河道场景专属defget_train_aug():returnA.Compose([A.Resize(640,640),# 统一图像尺寸适配航拍宽画幅A.HorizontalFlip(p0.5),# 随机左右翻转A.RandomBrightnessContrast(p0.3),# 明暗、对比度调整削弱水面反光A.Normalize(),ToTensorV2()])# 定义验证集数据增强仅归一化不做随机变换defget_val_aug():returnA.Compose([A.Resize(640,640),A.Normalize(),ToTensorV2()])3.5 U-Net 模型训练脚本train_unet_water.py选用水利分割场景常用的U-Net模型结合644张样本特性与河道航拍场景调优超参防止小数据集过拟合 无人机河道水体分割 - U-Net训练脚本 数据集规模644张航拍图像 适配场景水面反光、岸线遮挡、野外河道全域水体分割 核心优化小数据集防过拟合、水面特征强化 importtorchimporttorch.nnasnnfromtorch.utils.dataimportDataLoaderimportsegmentation_models_pytorchassmpfromtqdmimporttqdmfromwater_datasetimportWaterSegmentDataset,get_train_aug,get_val_aug# 全局配置河道场景专属参数DEVICEtorch.device(cudaiftorch.cuda.is_available()elsecpu)# 数据集路径TRAIN_IMG./water_segment_dataset/images/trainTRAIN_MASK./water_segment_dataset/masks/trainVAL_IMG./water_segment_dataset/images/valVAL_MASK./water_segment_dataset/masks/val# 训练超参644张小数据集降低batch、减少轮次防过拟合BATCH_SIZE8EPOCHS50LEARNING_RATE1e-4# deftrain():# 1. 加载训练集、验证集train_datasetWaterSegmentDataset(TRAIN_IMG,TRAIN_MASK,transformget_train_aug())val_datasetWaterSegmentDataset(VAL_IMG,VAL_MASK,transformget_val_aug())train_loaderDataLoader(train_dataset,batch_sizeBATCH_SIZE,shuffleTrue,num_workers2)val_loaderDataLoader(val_dataset,batch_sizeBATCH_SIZE,shuffleFalse,num_workers2)# 2. 初始化U-Net模型骨干网络选用efficientnet-b0轻量化适配边缘部署modelsmp.Unet(encoder_nameefficientnet-b0,encoder_weightsimagenet,# 使用预训练权重小数据集加速收敛in_channels3,classes1,# 输出单通道掩码二分类分割activationsigmoid).to(DEVICE)# 3. 损失函数 优化器水体分割二分类任务首选BCE损失loss_fnnn.BCELoss()optimizertorch.optim.Adam(model.parameters(),lrLEARNING_RATE)# 4. 开始迭代训练forepochinrange(EPOCHS):# 训练阶段model.train()train_loss0.0pbartqdm(train_loader,descfEpoch{epoch1}/{EPOCHS}Train)forimgs,masksinpbar:imgs,masksimgs.to(DEVICE),masks.to(DEVICE)optimizer.zero_grad()outputsmodel(imgs)lossloss_fn(outputs,masks.unsqueeze(1))loss.backward()optimizer.step()train_lossloss.item()pbar.set_postfix({loss:loss.item()})# 验证阶段model.eval()val_loss0.0withtorch.no_grad():pbar_valtqdm(val_loader,descfEpoch{epoch1}/{EPOCHS}Val)forimgs,masksinpbar_val:imgs,masksimgs.to(DEVICE),masks.to(DEVICE)outputsmodel(imgs)lossloss_fn(outputs,masks.unsqueeze(1))val_lossloss.item()# 打印每轮平均损失avg_train_losstrain_loss/len(train_loader)avg_val_lossval_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加载训练好的模型对无人机航拍图像做水体分割推理输出水体掩码与叠加效果图适配巡检业务可视化需求 河道水体分割推理脚本 功能单张/批量航拍图像水体提取、分割结果可视化 落地场景无人机实时巡检画面分析、洪涝区域离线统计 importtorchimportcv2importnumpyasnpimportosimportsegmentation_models_pytorchassmbfromPILimportImage# 全局配置DEVICEtorch.device(cudaiftorch.cuda.is_available()elsecpu)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_okTrue)# 加载模型modelsmb.Unet(encoder_nameefficientnet-b0,encoder_weightsNone,in_channels3,classes1,activationsigmoid).to(DEVICE)model.load_state_dict(torch.load(MODEL_PATH,map_locationDEVICE))model.eval()defpredict_single_image(img_path):# 读取航拍图像imagecv2.imread(img_path)h,wimage.shape[:2]# 图像预处理img_rgbcv2.cvtColor(image,cv2.COLOR_BGR2RGB)img_resizecv2.resize(img_rgb,(640,640))img_tensortorch.from_numpy(img_resize).permute(2,0,1).float()/255.0img_tensorimg_tensor.unsqueeze(0).to(DEVICE)# 模型推理阈值0.5区分水体与背景适配水面模糊边界withtorch.no_grad():outputmodel(img_tensor)pred_mask(output0.5).float().squeeze().cpu().numpy()# 掩码还原原图尺寸pred_maskcv2.resize(pred_mask,(w,h))# 生成蓝色掩码图层叠加到原图直观展示水体区域mask_colornp.zeros_like(image)mask_color[pred_mask1][255,0,0]# 蓝色标注水体# 原图掩码融合resultcv2.addWeighted(image,0.7,mask_color,0.3,0)returnresult# 批量处理测试图像if__name____main__:forimg_nameinos.listdir(TEST_IMG_DIR):ifimg_name.endswith((.jpg,.png)):img_full_pathos.path.join(TEST_IMG_DIR,img_name)seg_resultpredict_single_image(img_full_path)save_pathos.path.join(SAVE_RESULT_DIR,img_name)cv2.imwrite(save_path,seg_result)print(f批量水体分割完成结果已保存至{SAVE_RESULT_DIR})四、数据集使用策略 场景化深度思考4.1 数据划分建议数据集共644张航拍图像结合语义分割任务特性与小样本训练规则推荐划分比例训练集70%451张用于模型主体训练验证集15%97张用于训练过程中监控精度、防止过拟合测试集15%96张用于模型最终效果评估与业务仿真测试。划分时保证不同光照、不同河道形态样本均匀分布避免局部场景数据集中。4.2 河道航拍分割场景调优经验应对水面反光干扰水面强反光是河道分割最大难点训练阶段开启亮度、对比度随机变换推理时可搭配图像滤波预处理能明显提升边界分割精度。小样本防过拟合本数据集仅644张图像训练时建议选用轻量化骨干网络、降低BatchSize、配合早停机制不要使用过深的网络结构。边缘部署适配U-Net轻量化版本推理速度快训练后的模型可直接部署在无人机机载端、水利边缘网关实现航拍实时分割端到端巡检。边界优化河道水岸交界区域像素模糊可在后处理中使用形态学运算优化分割边缘提升水域面积统计精度。4.3 技术拓展方向多任务融合在水体分割基础上叠加目标检测同时实现水域分割河道漂浮物、违建船只识别打造一体化智能巡检系统。时序分析结合多时段航拍时序图像分析水体面积变化实现流域水位、水量动态监测。跨场景迁移基于本数据集做预训练再用本地流域数据微调快速适配不同地域河道、湖泊、水库场景。五、总结本套无人机河道航拍水体语义分割数据集聚焦水利行业刚需场景644张标注样本场景丰富、标注规范、格式通用完美适配语义分割算法训练与水利AI工程落地。配套U-Net全流程代码覆盖数据加载、训练、推理、可视化从算法研发到业务部署形成完整链路。该数据集既适合AI学习者入门语义分割任务也可作为水利、环保企业智慧巡检、洪涝监测项目的基础数据底座能够大幅降低水利计算机视觉项目的研发成本与周期。附加标签#U-Net实战 #语义分割开源 #智慧水利 #航拍图像AI分析 #洪涝区域识别
