尧图网站建设 
简介面向风力发电机叶片智能监测与语义分割研究的图像数据集与Python训练代码适合计算机视觉、智慧风电方向的学生和算法工程师用于叶片表面状态识别、磨损与裂缝检测等场景。压缩包共2000个文件主要包含约1994张tif格式的风扇叶片原始拍摄图和配套标签图另有6个Python脚本提供数据划分、图像预处理、U-Net网络构建、模型训练和预测等模块可直接运行并调整参数包体约810.93MB。已有112人学习下载。数据集覆盖不同拍摄环境和光照条件下的叶片样本标注相对完整有助于复现语义分割实验配合脚本中的数据集加载器和预处理流程可较方便地进行数据集替换与迁移训练适合希望快速验证分割算法、降低工程实现成本的研究与开发人员。1. 风力发电机风扇语义分割数据集拿来就能跑通 U-Net 吗做风电叶片巡检的手上从来不缺图缺的是标好标签的数据集和一套能直接改的训练代码。风力发电机风扇语义分割数据集正好补上这一块大量叶片图像、对应的像素级标签连同 Unet.py、train.py、dataset.py、predict.py、pre_process.py、split_data.py 等 Python 脚本从数据划分到模型训练再到推理出掩码整条链路都是通的。它适合两类人一类是做风电设备视觉检测的工程师想验证语义分割在叶片污损、裂纹识别上的效果另一类是刚接触语义分割的开发者想用现成的真实数据把 U-Net 完整跑一遍。下面按代码文件逐个拆讲清参数怎么设、哪些地方容易翻车。2. U-Net 与训练链路从 Unet.py 到 train.py 的参数拆解2.1 Unet.py 网络结构编码器、解码器与跳跃连接U-Net 在分割任务里属于“稳”字当头的结构Unet.py 实现的正是经典版本编码器逐层提取特征并压低分辨率解码器把低分辨率特征图逐步恢复回去中间通过跳跃连接把同尺度的编码器特征拼到解码器特征上。这样做的意义在于叶片边缘、裂纹末梢这类高分辨率细节不会因为下采样而彻底丢失。对风机叶片这种目标细长、边缘清晰的场景U-Net 的稳定性往往比直接套分类网络做像素预测好得多。Unet.py 里的 DoubleConv 是整个网络的骨架模块负责两次卷积加归一化import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super(DoubleConv, self).__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)参数含义in_ch 是输入通道数第一层对应图像 RGB 三通道值为 3out_ch 是本层输出的特征图通道数常见配置是 64、128、256 逐层翻倍。padding1 保证卷积不改变特征图尺寸方便后续和同尺度的跳跃连接拼接。BatchNorm 的作用是稳定中间层数据分布让训练收敛更快。编码器和解码器的配合是 U-Net 的核心结构上大致是这样class UNet(nn.Module): def __init__(self, in_channels3, num_classes2): super(UNet, self).__init__() self.inc DoubleConv(in_channels, 64) self.down1 nn.Sequential(nn.MaxPool2d(2), DoubleConv(64, 128)) self.down2 nn.Sequential(nn.MaxPool2d(2), DoubleConv(128, 256)) self.down3 nn.Sequential(nn.MaxPool2d(2), DoubleConv(256, 512)) self.up2 nn.ConvTranspose2d(512, 256, 2, stride2) self.conv2 DoubleConv(512, 256) self.up3 nn.ConvTranspose2d(256, 128, 2, stride2) self.conv3 DoubleConv(256, 128) self.up4 nn.ConvTranspose2d(128, 64, 2, stride2) self.conv4 DoubleConv(128, 64) self.outc nn.Conv2d(64, num_classes, 1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x self.up2(x4) x torch.cat([x, x3], dim1) x self.conv2(x) x self.up3(x) x torch.cat([x, x2], dim1) x self.conv3(x) x self.up4(x) x torch.cat([x, x1], dim1) x self.conv4(x) x self.outc(x) return x这段是 U-Net 的标准结构上采样完成后用 torch.cat 沿通道维度拼接编码器同尺度的特征再过一层 DoubleConv 把翻倍的通道压回去。实际 Unet.py 里一般还会把编码器加深到 1024 通道对小数据集保留三层就够参数量减少接近一半训练更快也不容易过拟合。需要改的地方就两个in_channels 对应图像通道通常就是 3num_classes 对应类别数。如果只做叶片/背景二分类num_classes 填 2多分类就填实际标注类别数。最后输出的特征图通道数等于 num_classes后面配合 CrossEntropyLoss 直接使用。注意跳跃连接拼接后一定要跟一层卷积把通道压回来否则解码器通道数会叠加失控模型参数成倍膨胀但精度不会提升。2.2 dataset.py 图像加载图像-标签配对与 transformdataset.py 里的 Dataset 类承担图像和标签的配对加载。训练时 DataLoader 每次取一个 batch就会调用一次__getitem__把图像张量和标签张量一起返回。这部分的写法直接决定数据会不会读错、标签会不会错位。import os import numpy as np import torch from PIL import Image from torch.utils.data import Dataset class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone, mask_transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform self.mask_transform mask_transform self.img_names sorted(os.listdir(img_dir)) def __len__(self): return len(self.img_names) def __getitem__(self, idx): img_name self.img_names[idx] image Image.open(os.path.join(self.img_dir, img_name)).convert(RGB) mask Image.open(os.path.join(self.mask_dir, img_name)).convert(L) if self.transform: image self.transform(image) if self.mask_transform: mask self.mask_transform(mask) else: mask torch.from_numpy(np.array(mask)).long() return image, mask这里的配对逻辑很简单但有一个前提图像文件名和标签文件名必须严格一一对应。这个数据集在整理时保持了命名一致性但如果是别的项目拿过来复用图像叫img_001.jpg、标签叫mask_001.png这类情况很常见你就需要在代码里做前缀统一而不是只靠sorted()按顺序硬对齐。一旦名字错位训练出来的模型等于用错误的监督信号结果自然会乱。mask_transform 单独设计是因为标签的变换和图像不同标签不能做 Normalize。我一般会拆成两个 transformfrom torchvision import transforms def get_transform(): img_transform transforms.Compose([ transforms.Resize((512, 512)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) mask_transform transforms.Compose([ transforms.Resize((512, 512), interpolationtransforms.InterpolationMode.NEAREST), transforms.ToTensor() ]) return img_transform, mask_transformResize 的参数在这里有讲究图像用默认的双线性插值没问题但标签必须指定 NEAREST否则 0/1 边界会被插值成灰色过渡像素这就是“标签被平滑”的典型来源一旦发生网络学到的边界就是模糊的。ToTensor 会把 PIL 图像的 HWC 转成 CHW并把图像像素归一化到 0~1对标签来说ToTensor 后得到 0.0 和 1.0 的 float 张量后面计算交叉熵时再转 long。2.3 train.py 训练参数loss、优化器、batch 与 epoch 怎么定train.py 把前面的模块串起来完成一个标准训练循环import torch import torch.nn as nn from torch.utils.data import DataLoader from dataset import SegDataset, get_transform from Unet import UNet img_transform, mask_transform get_transform() train_set SegDataset(DATASET/train/images, DATASET/train/masks, transformimg_transform, mask_transformmask_transform) train_loader DataLoader(train_set, batch_size8, shuffleTrue, num_workers4, pin_memoryTrue) model UNet(in_channels3, num_classes2) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-4) for epoch in range(100): model.train() total_loss 0.0 for images, masks in train_loader: images images.to(device) masks masks.to(device).squeeze(1).long() outputs model(images) loss criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() print(fepoch {epoch 1:03d} loss {total_loss / len(train_loader):.4f}) torch.save(model.state_dict(), fcheckpoints/epoch_{epoch 1:03d}.pth)逻辑说明每次迭代先清空梯度前向计算得到输出计算 loss反向传播再更新权重。masks 从 (B, 1, H, W) 压缩成 (B, H, W) 并转 long这是因为 CrossEntropyLoss 期望的 target 是整数索引而不是 one-hot 编码或 float 类型。这一步写错训练时会直接报错或者 loss 变成 NaN。关键训练参数的推荐值如下参数推荐值说明batch_size8512×512 输入时约占用 6~8G 显存显存小降到 2~4lr1e-4Adam 的稳妥起点过大容易震荡过小收敛慢epoch100配合早停看验证 IoU 选最优权重num_workers4数据读取瓶颈明显时调到 8num_classes2与数据集标注类别数一致训练时不要只会按 epoch 存权重更推荐按验证指标存最优权重best_iou 0 for epoch in range(100): ... iou validate(model, val_loader) if iou best_iou: best_iou iou torch.save(model.state_dict(), best_model.pth)这个best_model.pth就是后面 predict.py 要用的权重。这里有个习惯值得养成不按最后的 epoch 存而是按验证指标存最优状态否则最后几个 epoch 可能是过拟合状态直接拿 final 权重去预测反而效果差。3. 数据准备与增强pre_process.py 和 split_data.py 的正确用法3.1 预处理TIF 大图切片与归一化数据集里以 .tif 结尾的原始图分辨率通常远大于训练输入尺寸直接把整张图 Resize 到 512 会让叶片裂纹变成几个像素宽分割精度无从谈起。pre_process.py 做的事可以归纳为把大图切成固定大小的 patch、做归一化、附加数据增强。切片逻辑大致是这样import cv2 import numpy as np def slide_crop(image, crop_size512, step512): h, w image.shape[:2] patches [] for y in range(0, h - crop_size 1, step): for x in range(0, w - crop_size 1, step): patches.append(image[y:y crop_size, x:x crop_size]) return patchescrop_size 是每块 patch 的边长step 是滑窗步长。当 step 等于 crop_size 时patch 之间没有重叠当 step 小于 crop_size 时产生重叠相当于对位置做了增广。我一般把 step 设为 crop_size 的二分之一这样每张大图能衍生出多块带重叠的 patch训练样本量直接扩大。代价是相邻 patch 高度相似会增加训练时间和轻微过拟合风险数据量充足时这么用数据量少时更推荐。切片必须同步处理原图和标签这里最容易出错def crop_pair(image, mask, crop_size512, step512): h, w image.shape[:2] for y in range(0, h - crop_size 1, step): for x in range(0, w - crop_size 1, step): img_patch image[y:y crop_size, x:x crop_size] mask_patch mask[y:y crop_size, x:x crop_size] yield img_patch, mask_patch这个函数用生成器方式产出 patch 对原图和标签使用完全相同的坐标窗口裁剪。最怕的情况是原图裁了、标签没裁或者坐标偏移训练时 loss 完全不下降检查半天才发现标签对不上。归一化方面使用 ImageNet 统计量是通用做法前文代码里已经写过。这里要强调的是预处理必须统一训练和推理用同一套 mean/std。如果在训练时用 (0.485, 0.456, 0.406)推理时忘了 Normalize预测效果会直接劣化而且劣化得很隐蔽掩码边缘会出大量噪声。这类问题不好定位我通常会先在 predict.py 里打印输入张量的 mean 和 std确认它和训练时一致再往下排查。叶片分割经常会遇到光照干扰。如果发现模型在暗光下失效可以在预处理里加上光照扰动def adjust_brightness(image, factor1.2): hsv cv2.cvtColor(image, cv2.COLOR_BGR2HSV) hsv[:, :, 2] np.clip(hsv[:, :, 2] * factor, 0, 255) return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)这种增强要放在切片之后否则增强只作用于整张原图无法体现不同区域的光照差异。增强操作只作用于原图标签像素值保持不变。3.2 split_data.py数据集划分比例与随机种子split_data.py 负责把全量数据按比例拆成训练、验证、测试三个子集。核心逻辑如下import os import random random.seed(42) images sorted(os.listdir(DATASET/images)) n len(images) idx list(range(n)) random.shuffle(idx) train_ratio, val_ratio, test_ratio 0.7, 0.15, 0.15 train_idx idx[:int(n * train_ratio)] val_idx idx[int(n * train_ratio):int(n * (train_ratio val_ratio))] test_idx idx[int(n * (train_ratio val_ratio)):] train_imgs [images[i] for i in train_idx] val_imgs [images[i] for i in val_idx] test_imgs [images[i] for i in test_idx]train 0.7、val 0.15、test 0.15 是常见默认值。如果数据总量只有几百张建议 train 提到 0.8test 保持 0.1因为分割任务需要足够多的样本覆盖不同拍摄角度和光照条件val 比例过高反而浪费训练样本。random.seed(42)这句经常被忽略但不写的话每次运行划分结果都不同前后实验无法对比调参就像在抓瞎。固定种子后训练集不变模型改进带来的指标提升才是真实提升。我通常还会把划分后的名单导出成文件with open(train.txt, w) as f: f.write(\n.join([os.path.basename(p) for p in train_imgs]))这份文件有双重用途一是下次训练直接读名单不重复执行划分二是复核标签时手动抽查某个文件名对应的图像。配合 dataset.py 里按名字配对的逻辑只要名单一致图像和标签就不会错位。划分时还要注意按场景或批次分组。随机打乱看起来公平但同一场景的连续帧往往高度相似如果它们分别落在 train 和 val验证 IoU 会被虚高真实泛化能力被高估。干净的做法是先把文件名按场景前缀分组以组为单位划分。这在风电巡检场景里尤其重要因为无人机一次起飞的连续照片都在同一个光照条件下。3.3 划分后的自检标签像素值、尺寸和通道划分完成的标准不是目录建好了而是脚本检查通过。我每次都会写一段自检代码在训练前把数据集完整过一遍from PIL import Image import numpy as np import os for subset in [train, val, test]: img_dir fDATASET/{subset}/images mask_dir fDATASET/{subset}/masks for name in os.listdir(img_dir): img np.array(Image.open(os.path.join(img_dir, name)).convert(RGB)) mask np.array(Image.open(os.path.join(mask_dir, name)).convert(L)) assert img.shape[0] mask.shape[0] and img.shape[1] mask.shape[1], \ f{name} size mismatch: {img.shape} vs {mask.shape} unique np.unique(mask) assert set(unique) {0, 1}, f{name} unexpected mask values: {unique}这段脚本做了两件事检查图像和标签尺寸是否一致确认标签像素值是否只含 0 和 1。如果标签里出现 255说明是位图格式训练前需要对标签做二值化如果标签是三通道的伪彩色图convert(L)会强制压成单通道但要注意颜色映射关系不能想当然。实际项目里标签格式五花八门。有的标注工具导出三通道 PNG看起来是黑白实际 shape 是 (H, W, 3)有的标签边缘有抗锯齿过渡值比如 0、128、255 三个值同时出现这类标签多数是标注软件的画笔边缘产生的建议先做阈值二值化小于 128 归 0大于等于 128 归 1再送训练。自检脚本跑完后还要随机抽几张图人工过一遍。把原图和标签叠加成半透明图确认叶片轮廓和标签大体一致。这一步虽然费时间但能发现标注错位、标注遗漏这类程序检查发现不了的问题。数据质量不过关后面模型怎么调参都是白费这是语义分割项目里最扎实的一条经验。4. predict.py 推理与常见问题排查从权重加载到掩码输出4.1 推理流程加载权重、推理与保存掩码predict.py 的作用是把训练好的 best_model.pth 加载进来对一张新图像做前向推理输出像素级掩码。完整流程分三步构建与训练时一致的模型、加载权重、对输入做同样的预处理再推理。import torch import numpy as np from PIL import Image from Unet import UNet from torchvision import transforms model UNet(in_channels3, num_classes2) ckpt torch.load(best_model.pth, map_locationcpu) model.load_state_dict(ckpt) model.eval() img Image.open(test.jpg).convert(RGB) transform transforms.Compose([ transforms.Resize((512, 512)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img_tensor transform(img).unsqueeze(0) with torch.no_grad(): output model(img_tensor) pred torch.argmax(output, dim1)load_state_dict之前必须保证模型结构和训练时完全一致包括 in_channels 和 num_classes。最常见的报错是size mismatch原因是训练时 num_classes3推理时建模型用了 num_classes2最后一层卷积权重形状对不上。最省心的做法是把超参数一起存进权重文件torch.save({state_dict: model.state_dict(), num_classes: num_classes}, best_model.pth)推理时先读 num_classes 再建模型整个流程就不容易犯低级错误。这个设计看起来只多写了几行但在换模型、换数据集时会省下很多调试时间。argmax(dim1)是在每个像素位置取最大值的类别索引dim1 对应 num_classes 那一维。如果 output 形状是 (1, 2, 512, 512)argmax 后得到 (1, 512, 512)0 是背景1 是叶片。保存掩码时转成 uint8 再乘 255图像查看器才能正常显示为黑底白叶pred pred.cpu().numpy().astype(np.uint8) Image.fromarray(pred[0] * 255).save(result_mask.png)如果训练时设置了多分类比如 0 背景、1 正常叶片、2 裂纹、3 污渍建议给每个类别指定不同灰度值保存否则肉眼很难区分。预测的类别编号要和训练标签编号严格对应这条在 split_data 阶段的自检代码里就要确认好。4.2 五个高频坑现象、原因、解决下面这五个问题是我在跑这类数据集时反复遇到的每一条都是按“现象 → 原因 → 解决”的顺序记录。坑一CUDA out of memory。现象DataLoader 传入第一个 batch 后进程直接报 CUDA out of memory 退出。 原因512×512 输入、batch_size8、U-Net 中间多通道特征图累计显存超过 8G。 解决优先把 batch_size 从 8 降到 4显存占用几乎减半再不够就把 Resize 从 512 降到 256。另一个通用手段是加混合精度训练把前向计算和 loss 计算包在torch.cuda.amp.autocast()里能省 30% 左右的显存几乎不影响精度。坑二loss 一直维持在 0.69 上下不动。现象训练几十个 epochloss 在一个固定值附近震荡完全没有下降趋势。 原因0.69 对应二分类交叉熵的随机猜测概率说明网络没有学到有效特征。最常见的是标签全部为 0或标签是 255 导致 loss 计算异常。另一种常见原因是增强时对标签用了双线性插值把 0/1 边缘插成了小数。 解决先跑 3.3 节的自检脚本检查标签唯一值和形状标签是 0/255 就先二值化增强代码里 mask 的插值方式改成 NEAREST。遇到 loss 不降别先往玄学方向想九成是数据问题。坑三预测结果全黑或全白。现象推理出的掩码所有像素都是同一个值完全看不出叶片轮廓。 原因网络未收敛或者 argmax 用错维度导致类别索引错位。如果 output 的形状和预期不一致说明 num_classes 设置错了。 解决先print(output.shape)确认输出形状再检查 predict.py 的 num_classes 和训练时是否一致最后确认加载的权重是不是按验证指标保存的最佳权重。这个排查顺序能覆盖绝大多数情况。坑四训练时 GPU 利用率上不去速度很慢。现象每个 epoch 耗时很长nvidia-smi 显示 GPU 利用率的实心区域经常断开。 原因DataLoader 的 num_workers0图像读取和预处理全在主进程串行执行图像解码成为瓶颈。 解决把 num_workers 调到 4 或 8并加上pin_memoryTrue。如果机器内存不够可以先用 pre_process.py 把大图全部切成小 patch 并缓存训练时直接读已切好的图I/O 压力会小很多。坑五验证指标虚高换成真实场景明显下降。现象验证集上 IoU 很不错但拿到新拍摄的图片上预测效果明显变差。 原因随机划分导致同一场景的连续帧同时泄漏到 train 和 val网络记住了场景特征而不是叶片本身。 解决按拍摄批次或场景分组后再划分数据集。风电巡检的无人机图片连续性很强这个坑尤其值得重视。我在做视觉项目时被这一条坑过很多次一旦指标和实际效果对不上第一件事永远是查数据划分逻辑。4.3 验证指标IoU 与 Dice 的快速计算分割任务的验证指标不能只看 loss。loss 是训练信号和实际效果并不完全正相关IoU 才是更接近业务结果的指标。计算代码通常是这样的def compute_iou(pred_mask, gt_mask, num_classes2): iou_list [] for cls in range(num_classes): pred_cls (pred_mask cls) gt_cls (gt_mask cls) intersection (pred_cls gt_cls).sum() union (pred_cls | gt_cls).sum() iou_list.append(intersection / (union 1e-6)) return iou_list分母加 1e-6 是为了避免除零。对叶片分割我只看前景类 IoU也就是类别 1 的 IoU背景类占比太大数值虚高没有参考价值。如果前景 IoU 能达到 0.85 以上模型基本具备实用价值0.7 以下则基本上只能处理与训练集高度相似的场景。Dice 系数也会一起算def dice_score(pred_mask, gt_mask, smooth1e-6): pred_flat pred_mask.reshape(-1) gt_flat gt_mask.reshape(-1) intersection (pred_flat * gt_flat).sum() return (2.0 * intersection smooth) / (pred_flat.sum() gt_flat.sum() smooth)Dice 对前景类别的敏感度和 IoU 略有不同类别不平衡时 Dice 能更敏锐地反映小目标的分割质量。建议两个都算结合来看。一个背景占 90% 以上的数据集里准确率 99% 都可能是全预测背景得到的所以只报告准确率的论文或者项目你都要留个心眼。5. 进阶把预测掩码变成叶片面积占比与故障统计5.1 从掩码统计叶片面积占比推理得到掩码后直接丢给业务方看意义不大需要把它转成可量化的指标。最常用的是叶片面积占比import numpy as np from PIL import Image mask np.array(Image.open(result_mask.png).convert(L)) leaf_ratio (mask 0).sum() / mask.size print(fleaf area ratio: {leaf_ratio:.4f})这个数值能用来快速判断当前画面里有没有叶片如果掩码几乎没有前景像素说明这一帧是纯天空背景可以跳过分帧处理如果占比突然偏离正常区间说明叶片姿态或遮挡出现异常。不同机型、不同拍摄距离下叶片在画面中的合理占比区间不一样可以先统计一批正常样本做个区间偏离区间就触发复核运维上比较实用。5.2 可视化叠加与结果输出光有面积占比还不够业务人员更希望看到“具体哪一块有问题”。把预测区域叠加到原图上输出是最直观的交付形式import cv2 img cv2.imread(test.jpg) overlay img.copy() overlay[mask 0] (0, 0, 255) result cv2.addWeighted(img, 0.6, overlay, 0.4, 0) cv2.imwrite(result_overlay.jpg, result)mask 是 (H, W) 的布尔数组覆盖到 overlay 上。原图 0.6、叠加层 0.4 是合适的透明比例既能看清叶片轮廓又不会完全遮住原图细节。多分类情况下给不同类别分配不同颜色比如裂纹标红、污渍标黄巡检人员打开图就能定位。验证方法也有讲究。拿到一个新批次的图片我通常先跑 20 张把每张的预测掩码和原图叠加后人工过一遍。重点看两个地方一是叶片边缘有没有连续断裂二是背景有没有被误判成叶片。前者说明标注边缘质量差或者模型容量不足后者说明训练数据里背景复杂度不够需要补充负样本。第二轮再全量跑结合面积占比做统计报表。从那以后我每次跑完一个新模型都强制走一遍这套流程先算前景 IoU 和 Dice再抽 20 张图做人工肉眼验证最后输出叠加图和面积占比。顺序不能反指标先发现问题叠加图再定位问题。希望帮到你。本文还有配套的精品资源点击获取 风力发电机叶片语义分割：U-Net数据集与训练代码全解析
简介面向风力发电机叶片智能监测与语义分割研究的图像数据集与Python训练代码适合计算机视觉、智慧风电方向的学生和算法工程师用于叶片表面状态识别、磨损与裂缝检测等场景。压缩包共2000个文件主要包含约1994张tif格式的风扇叶片原始拍摄图和配套标签图另有6个Python脚本提供数据划分、图像预处理、U-Net网络构建、模型训练和预测等模块可直接运行并调整参数包体约810.93MB。已有112人学习下载。数据集覆盖不同拍摄环境和光照条件下的叶片样本标注相对完整有助于复现语义分割实验配合脚本中的数据集加载器和预处理流程可较方便地进行数据集替换与迁移训练适合希望快速验证分割算法、降低工程实现成本的研究与开发人员。1. 风力发电机风扇语义分割数据集拿来就能跑通 U-Net 吗做风电叶片巡检的手上从来不缺图缺的是标好标签的数据集和一套能直接改的训练代码。风力发电机风扇语义分割数据集正好补上这一块大量叶片图像、对应的像素级标签连同 Unet.py、train.py、dataset.py、predict.py、pre_process.py、split_data.py 等 Python 脚本从数据划分到模型训练再到推理出掩码整条链路都是通的。它适合两类人一类是做风电设备视觉检测的工程师想验证语义分割在叶片污损、裂纹识别上的效果另一类是刚接触语义分割的开发者想用现成的真实数据把 U-Net 完整跑一遍。下面按代码文件逐个拆讲清参数怎么设、哪些地方容易翻车。2. U-Net 与训练链路从 Unet.py 到 train.py 的参数拆解2.1 Unet.py 网络结构编码器、解码器与跳跃连接U-Net 在分割任务里属于“稳”字当头的结构Unet.py 实现的正是经典版本编码器逐层提取特征并压低分辨率解码器把低分辨率特征图逐步恢复回去中间通过跳跃连接把同尺度的编码器特征拼到解码器特征上。这样做的意义在于叶片边缘、裂纹末梢这类高分辨率细节不会因为下采样而彻底丢失。对风机叶片这种目标细长、边缘清晰的场景U-Net 的稳定性往往比直接套分类网络做像素预测好得多。Unet.py 里的 DoubleConv 是整个网络的骨架模块负责两次卷积加归一化import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super(DoubleConv, self).__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x)参数含义in_ch 是输入通道数第一层对应图像 RGB 三通道值为 3out_ch 是本层输出的特征图通道数常见配置是 64、128、256 逐层翻倍。padding1 保证卷积不改变特征图尺寸方便后续和同尺度的跳跃连接拼接。BatchNorm 的作用是稳定中间层数据分布让训练收敛更快。编码器和解码器的配合是 U-Net 的核心结构上大致是这样class UNet(nn.Module): def __init__(self, in_channels3, num_classes2): super(UNet, self).__init__() self.inc DoubleConv(in_channels, 64) self.down1 nn.Sequential(nn.MaxPool2d(2), DoubleConv(64, 128)) self.down2 nn.Sequential(nn.MaxPool2d(2), DoubleConv(128, 256)) self.down3 nn.Sequential(nn.MaxPool2d(2), DoubleConv(256, 512)) self.up2 nn.ConvTranspose2d(512, 256, 2, stride2) self.conv2 DoubleConv(512, 256) self.up3 nn.ConvTranspose2d(256, 128, 2, stride2) self.conv3 DoubleConv(256, 128) self.up4 nn.ConvTranspose2d(128, 64, 2, stride2) self.conv4 DoubleConv(128, 64) self.outc nn.Conv2d(64, num_classes, 1) def forward(self, x): x1 self.inc(x) x2 self.down1(x1) x3 self.down2(x2) x4 self.down3(x3) x self.up2(x4) x torch.cat([x, x3], dim1) x self.conv2(x) x self.up3(x) x torch.cat([x, x2], dim1) x self.conv3(x) x self.up4(x) x torch.cat([x, x1], dim1) x self.conv4(x) x self.outc(x) return x这段是 U-Net 的标准结构上采样完成后用 torch.cat 沿通道维度拼接编码器同尺度的特征再过一层 DoubleConv 把翻倍的通道压回去。实际 Unet.py 里一般还会把编码器加深到 1024 通道对小数据集保留三层就够参数量减少接近一半训练更快也不容易过拟合。需要改的地方就两个in_channels 对应图像通道通常就是 3num_classes 对应类别数。如果只做叶片/背景二分类num_classes 填 2多分类就填实际标注类别数。最后输出的特征图通道数等于 num_classes后面配合 CrossEntropyLoss 直接使用。注意跳跃连接拼接后一定要跟一层卷积把通道压回来否则解码器通道数会叠加失控模型参数成倍膨胀但精度不会提升。2.2 dataset.py 图像加载图像-标签配对与 transformdataset.py 里的 Dataset 类承担图像和标签的配对加载。训练时 DataLoader 每次取一个 batch就会调用一次__getitem__把图像张量和标签张量一起返回。这部分的写法直接决定数据会不会读错、标签会不会错位。import os import numpy as np import torch from PIL import Image from torch.utils.data import Dataset class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone, mask_transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform self.mask_transform mask_transform self.img_names sorted(os.listdir(img_dir)) def __len__(self): return len(self.img_names) def __getitem__(self, idx): img_name self.img_names[idx] image Image.open(os.path.join(self.img_dir, img_name)).convert(RGB) mask Image.open(os.path.join(self.mask_dir, img_name)).convert(L) if self.transform: image self.transform(image) if self.mask_transform: mask self.mask_transform(mask) else: mask torch.from_numpy(np.array(mask)).long() return image, mask这里的配对逻辑很简单但有一个前提图像文件名和标签文件名必须严格一一对应。这个数据集在整理时保持了命名一致性但如果是别的项目拿过来复用图像叫img_001.jpg、标签叫mask_001.png这类情况很常见你就需要在代码里做前缀统一而不是只靠sorted()按顺序硬对齐。一旦名字错位训练出来的模型等于用错误的监督信号结果自然会乱。mask_transform 单独设计是因为标签的变换和图像不同标签不能做 Normalize。我一般会拆成两个 transformfrom torchvision import transforms def get_transform(): img_transform transforms.Compose([ transforms.Resize((512, 512)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) mask_transform transforms.Compose([ transforms.Resize((512, 512), interpolationtransforms.InterpolationMode.NEAREST), transforms.ToTensor() ]) return img_transform, mask_transformResize 的参数在这里有讲究图像用默认的双线性插值没问题但标签必须指定 NEAREST否则 0/1 边界会被插值成灰色过渡像素这就是“标签被平滑”的典型来源一旦发生网络学到的边界就是模糊的。ToTensor 会把 PIL 图像的 HWC 转成 CHW并把图像像素归一化到 0~1对标签来说ToTensor 后得到 0.0 和 1.0 的 float 张量后面计算交叉熵时再转 long。2.3 train.py 训练参数loss、优化器、batch 与 epoch 怎么定train.py 把前面的模块串起来完成一个标准训练循环import torch import torch.nn as nn from torch.utils.data import DataLoader from dataset import SegDataset, get_transform from Unet import UNet img_transform, mask_transform get_transform() train_set SegDataset(DATASET/train/images, DATASET/train/masks, transformimg_transform, mask_transformmask_transform) train_loader DataLoader(train_set, batch_size8, shuffleTrue, num_workers4, pin_memoryTrue) model UNet(in_channels3, num_classes2) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-4) for epoch in range(100): model.train() total_loss 0.0 for images, masks in train_loader: images images.to(device) masks masks.to(device).squeeze(1).long() outputs model(images) loss criterion(outputs, masks) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() print(fepoch {epoch 1:03d} loss {total_loss / len(train_loader):.4f}) torch.save(model.state_dict(), fcheckpoints/epoch_{epoch 1:03d}.pth)逻辑说明每次迭代先清空梯度前向计算得到输出计算 loss反向传播再更新权重。masks 从 (B, 1, H, W) 压缩成 (B, H, W) 并转 long这是因为 CrossEntropyLoss 期望的 target 是整数索引而不是 one-hot 编码或 float 类型。这一步写错训练时会直接报错或者 loss 变成 NaN。关键训练参数的推荐值如下参数推荐值说明batch_size8512×512 输入时约占用 6~8G 显存显存小降到 2~4lr1e-4Adam 的稳妥起点过大容易震荡过小收敛慢epoch100配合早停看验证 IoU 选最优权重num_workers4数据读取瓶颈明显时调到 8num_classes2与数据集标注类别数一致训练时不要只会按 epoch 存权重更推荐按验证指标存最优权重best_iou 0 for epoch in range(100): ... iou validate(model, val_loader) if iou best_iou: best_iou iou torch.save(model.state_dict(), best_model.pth)这个best_model.pth就是后面 predict.py 要用的权重。这里有个习惯值得养成不按最后的 epoch 存而是按验证指标存最优状态否则最后几个 epoch 可能是过拟合状态直接拿 final 权重去预测反而效果差。3. 数据准备与增强pre_process.py 和 split_data.py 的正确用法3.1 预处理TIF 大图切片与归一化数据集里以 .tif 结尾的原始图分辨率通常远大于训练输入尺寸直接把整张图 Resize 到 512 会让叶片裂纹变成几个像素宽分割精度无从谈起。pre_process.py 做的事可以归纳为把大图切成固定大小的 patch、做归一化、附加数据增强。切片逻辑大致是这样import cv2 import numpy as np def slide_crop(image, crop_size512, step512): h, w image.shape[:2] patches [] for y in range(0, h - crop_size 1, step): for x in range(0, w - crop_size 1, step): patches.append(image[y:y crop_size, x:x crop_size]) return patchescrop_size 是每块 patch 的边长step 是滑窗步长。当 step 等于 crop_size 时patch 之间没有重叠当 step 小于 crop_size 时产生重叠相当于对位置做了增广。我一般把 step 设为 crop_size 的二分之一这样每张大图能衍生出多块带重叠的 patch训练样本量直接扩大。代价是相邻 patch 高度相似会增加训练时间和轻微过拟合风险数据量充足时这么用数据量少时更推荐。切片必须同步处理原图和标签这里最容易出错def crop_pair(image, mask, crop_size512, step512): h, w image.shape[:2] for y in range(0, h - crop_size 1, step): for x in range(0, w - crop_size 1, step): img_patch image[y:y crop_size, x:x crop_size] mask_patch mask[y:y crop_size, x:x crop_size] yield img_patch, mask_patch这个函数用生成器方式产出 patch 对原图和标签使用完全相同的坐标窗口裁剪。最怕的情况是原图裁了、标签没裁或者坐标偏移训练时 loss 完全不下降检查半天才发现标签对不上。归一化方面使用 ImageNet 统计量是通用做法前文代码里已经写过。这里要强调的是预处理必须统一训练和推理用同一套 mean/std。如果在训练时用 (0.485, 0.456, 0.406)推理时忘了 Normalize预测效果会直接劣化而且劣化得很隐蔽掩码边缘会出大量噪声。这类问题不好定位我通常会先在 predict.py 里打印输入张量的 mean 和 std确认它和训练时一致再往下排查。叶片分割经常会遇到光照干扰。如果发现模型在暗光下失效可以在预处理里加上光照扰动def adjust_brightness(image, factor1.2): hsv cv2.cvtColor(image, cv2.COLOR_BGR2HSV) hsv[:, :, 2] np.clip(hsv[:, :, 2] * factor, 0, 255) return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)这种增强要放在切片之后否则增强只作用于整张原图无法体现不同区域的光照差异。增强操作只作用于原图标签像素值保持不变。3.2 split_data.py数据集划分比例与随机种子split_data.py 负责把全量数据按比例拆成训练、验证、测试三个子集。核心逻辑如下import os import random random.seed(42) images sorted(os.listdir(DATASET/images)) n len(images) idx list(range(n)) random.shuffle(idx) train_ratio, val_ratio, test_ratio 0.7, 0.15, 0.15 train_idx idx[:int(n * train_ratio)] val_idx idx[int(n * train_ratio):int(n * (train_ratio val_ratio))] test_idx idx[int(n * (train_ratio val_ratio)):] train_imgs [images[i] for i in train_idx] val_imgs [images[i] for i in val_idx] test_imgs [images[i] for i in test_idx]train 0.7、val 0.15、test 0.15 是常见默认值。如果数据总量只有几百张建议 train 提到 0.8test 保持 0.1因为分割任务需要足够多的样本覆盖不同拍摄角度和光照条件val 比例过高反而浪费训练样本。random.seed(42)这句经常被忽略但不写的话每次运行划分结果都不同前后实验无法对比调参就像在抓瞎。固定种子后训练集不变模型改进带来的指标提升才是真实提升。我通常还会把划分后的名单导出成文件with open(train.txt, w) as f: f.write(\n.join([os.path.basename(p) for p in train_imgs]))这份文件有双重用途一是下次训练直接读名单不重复执行划分二是复核标签时手动抽查某个文件名对应的图像。配合 dataset.py 里按名字配对的逻辑只要名单一致图像和标签就不会错位。划分时还要注意按场景或批次分组。随机打乱看起来公平但同一场景的连续帧往往高度相似如果它们分别落在 train 和 val验证 IoU 会被虚高真实泛化能力被高估。干净的做法是先把文件名按场景前缀分组以组为单位划分。这在风电巡检场景里尤其重要因为无人机一次起飞的连续照片都在同一个光照条件下。3.3 划分后的自检标签像素值、尺寸和通道划分完成的标准不是目录建好了而是脚本检查通过。我每次都会写一段自检代码在训练前把数据集完整过一遍from PIL import Image import numpy as np import os for subset in [train, val, test]: img_dir fDATASET/{subset}/images mask_dir fDATASET/{subset}/masks for name in os.listdir(img_dir): img np.array(Image.open(os.path.join(img_dir, name)).convert(RGB)) mask np.array(Image.open(os.path.join(mask_dir, name)).convert(L)) assert img.shape[0] mask.shape[0] and img.shape[1] mask.shape[1], \ f{name} size mismatch: {img.shape} vs {mask.shape} unique np.unique(mask) assert set(unique) {0, 1}, f{name} unexpected mask values: {unique}这段脚本做了两件事检查图像和标签尺寸是否一致确认标签像素值是否只含 0 和 1。如果标签里出现 255说明是位图格式训练前需要对标签做二值化如果标签是三通道的伪彩色图convert(L)会强制压成单通道但要注意颜色映射关系不能想当然。实际项目里标签格式五花八门。有的标注工具导出三通道 PNG看起来是黑白实际 shape 是 (H, W, 3)有的标签边缘有抗锯齿过渡值比如 0、128、255 三个值同时出现这类标签多数是标注软件的画笔边缘产生的建议先做阈值二值化小于 128 归 0大于等于 128 归 1再送训练。自检脚本跑完后还要随机抽几张图人工过一遍。把原图和标签叠加成半透明图确认叶片轮廓和标签大体一致。这一步虽然费时间但能发现标注错位、标注遗漏这类程序检查发现不了的问题。数据质量不过关后面模型怎么调参都是白费这是语义分割项目里最扎实的一条经验。4. predict.py 推理与常见问题排查从权重加载到掩码输出4.1 推理流程加载权重、推理与保存掩码predict.py 的作用是把训练好的 best_model.pth 加载进来对一张新图像做前向推理输出像素级掩码。完整流程分三步构建与训练时一致的模型、加载权重、对输入做同样的预处理再推理。import torch import numpy as np from PIL import Image from Unet import UNet from torchvision import transforms model UNet(in_channels3, num_classes2) ckpt torch.load(best_model.pth, map_locationcpu) model.load_state_dict(ckpt) model.eval() img Image.open(test.jpg).convert(RGB) transform transforms.Compose([ transforms.Resize((512, 512)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img_tensor transform(img).unsqueeze(0) with torch.no_grad(): output model(img_tensor) pred torch.argmax(output, dim1)load_state_dict之前必须保证模型结构和训练时完全一致包括 in_channels 和 num_classes。最常见的报错是size mismatch原因是训练时 num_classes3推理时建模型用了 num_classes2最后一层卷积权重形状对不上。最省心的做法是把超参数一起存进权重文件torch.save({state_dict: model.state_dict(), num_classes: num_classes}, best_model.pth)推理时先读 num_classes 再建模型整个流程就不容易犯低级错误。这个设计看起来只多写了几行但在换模型、换数据集时会省下很多调试时间。argmax(dim1)是在每个像素位置取最大值的类别索引dim1 对应 num_classes 那一维。如果 output 形状是 (1, 2, 512, 512)argmax 后得到 (1, 512, 512)0 是背景1 是叶片。保存掩码时转成 uint8 再乘 255图像查看器才能正常显示为黑底白叶pred pred.cpu().numpy().astype(np.uint8) Image.fromarray(pred[0] * 255).save(result_mask.png)如果训练时设置了多分类比如 0 背景、1 正常叶片、2 裂纹、3 污渍建议给每个类别指定不同灰度值保存否则肉眼很难区分。预测的类别编号要和训练标签编号严格对应这条在 split_data 阶段的自检代码里就要确认好。4.2 五个高频坑现象、原因、解决下面这五个问题是我在跑这类数据集时反复遇到的每一条都是按“现象 → 原因 → 解决”的顺序记录。坑一CUDA out of memory。现象DataLoader 传入第一个 batch 后进程直接报 CUDA out of memory 退出。 原因512×512 输入、batch_size8、U-Net 中间多通道特征图累计显存超过 8G。 解决优先把 batch_size 从 8 降到 4显存占用几乎减半再不够就把 Resize 从 512 降到 256。另一个通用手段是加混合精度训练把前向计算和 loss 计算包在torch.cuda.amp.autocast()里能省 30% 左右的显存几乎不影响精度。坑二loss 一直维持在 0.69 上下不动。现象训练几十个 epochloss 在一个固定值附近震荡完全没有下降趋势。 原因0.69 对应二分类交叉熵的随机猜测概率说明网络没有学到有效特征。最常见的是标签全部为 0或标签是 255 导致 loss 计算异常。另一种常见原因是增强时对标签用了双线性插值把 0/1 边缘插成了小数。 解决先跑 3.3 节的自检脚本检查标签唯一值和形状标签是 0/255 就先二值化增强代码里 mask 的插值方式改成 NEAREST。遇到 loss 不降别先往玄学方向想九成是数据问题。坑三预测结果全黑或全白。现象推理出的掩码所有像素都是同一个值完全看不出叶片轮廓。 原因网络未收敛或者 argmax 用错维度导致类别索引错位。如果 output 的形状和预期不一致说明 num_classes 设置错了。 解决先print(output.shape)确认输出形状再检查 predict.py 的 num_classes 和训练时是否一致最后确认加载的权重是不是按验证指标保存的最佳权重。这个排查顺序能覆盖绝大多数情况。坑四训练时 GPU 利用率上不去速度很慢。现象每个 epoch 耗时很长nvidia-smi 显示 GPU 利用率的实心区域经常断开。 原因DataLoader 的 num_workers0图像读取和预处理全在主进程串行执行图像解码成为瓶颈。 解决把 num_workers 调到 4 或 8并加上pin_memoryTrue。如果机器内存不够可以先用 pre_process.py 把大图全部切成小 patch 并缓存训练时直接读已切好的图I/O 压力会小很多。坑五验证指标虚高换成真实场景明显下降。现象验证集上 IoU 很不错但拿到新拍摄的图片上预测效果明显变差。 原因随机划分导致同一场景的连续帧同时泄漏到 train 和 val网络记住了场景特征而不是叶片本身。 解决按拍摄批次或场景分组后再划分数据集。风电巡检的无人机图片连续性很强这个坑尤其值得重视。我在做视觉项目时被这一条坑过很多次一旦指标和实际效果对不上第一件事永远是查数据划分逻辑。4.3 验证指标IoU 与 Dice 的快速计算分割任务的验证指标不能只看 loss。loss 是训练信号和实际效果并不完全正相关IoU 才是更接近业务结果的指标。计算代码通常是这样的def compute_iou(pred_mask, gt_mask, num_classes2): iou_list [] for cls in range(num_classes): pred_cls (pred_mask cls) gt_cls (gt_mask cls) intersection (pred_cls gt_cls).sum() union (pred_cls | gt_cls).sum() iou_list.append(intersection / (union 1e-6)) return iou_list分母加 1e-6 是为了避免除零。对叶片分割我只看前景类 IoU也就是类别 1 的 IoU背景类占比太大数值虚高没有参考价值。如果前景 IoU 能达到 0.85 以上模型基本具备实用价值0.7 以下则基本上只能处理与训练集高度相似的场景。Dice 系数也会一起算def dice_score(pred_mask, gt_mask, smooth1e-6): pred_flat pred_mask.reshape(-1) gt_flat gt_mask.reshape(-1) intersection (pred_flat * gt_flat).sum() return (2.0 * intersection smooth) / (pred_flat.sum() gt_flat.sum() smooth)Dice 对前景类别的敏感度和 IoU 略有不同类别不平衡时 Dice 能更敏锐地反映小目标的分割质量。建议两个都算结合来看。一个背景占 90% 以上的数据集里准确率 99% 都可能是全预测背景得到的所以只报告准确率的论文或者项目你都要留个心眼。5. 进阶把预测掩码变成叶片面积占比与故障统计5.1 从掩码统计叶片面积占比推理得到掩码后直接丢给业务方看意义不大需要把它转成可量化的指标。最常用的是叶片面积占比import numpy as np from PIL import Image mask np.array(Image.open(result_mask.png).convert(L)) leaf_ratio (mask 0).sum() / mask.size print(fleaf area ratio: {leaf_ratio:.4f})这个数值能用来快速判断当前画面里有没有叶片如果掩码几乎没有前景像素说明这一帧是纯天空背景可以跳过分帧处理如果占比突然偏离正常区间说明叶片姿态或遮挡出现异常。不同机型、不同拍摄距离下叶片在画面中的合理占比区间不一样可以先统计一批正常样本做个区间偏离区间就触发复核运维上比较实用。5.2 可视化叠加与结果输出光有面积占比还不够业务人员更希望看到“具体哪一块有问题”。把预测区域叠加到原图上输出是最直观的交付形式import cv2 img cv2.imread(test.jpg) overlay img.copy() overlay[mask 0] (0, 0, 255) result cv2.addWeighted(img, 0.6, overlay, 0.4, 0) cv2.imwrite(result_overlay.jpg, result)mask 是 (H, W) 的布尔数组覆盖到 overlay 上。原图 0.6、叠加层 0.4 是合适的透明比例既能看清叶片轮廓又不会完全遮住原图细节。多分类情况下给不同类别分配不同颜色比如裂纹标红、污渍标黄巡检人员打开图就能定位。验证方法也有讲究。拿到一个新批次的图片我通常先跑 20 张把每张的预测掩码和原图叠加后人工过一遍。重点看两个地方一是叶片边缘有没有连续断裂二是背景有没有被误判成叶片。前者说明标注边缘质量差或者模型容量不足后者说明训练数据里背景复杂度不够需要补充负样本。第二轮再全量跑结合面积占比做统计报表。从那以后我每次跑完一个新模型都强制走一遍这套流程先算前景 IoU 和 Dice再抽 20 张图做人工肉眼验证最后输出叠加图和面积占比。顺序不能反指标先发现问题叠加图再定位问题。希望帮到你。本文还有配套的精品资源点击获取