ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

风力发电机风扇语义分割:数据集与Python训练代码实战

风力发电机风扇语义分割:数据集与Python训练代码实战 简介这份资源面向计算机视觉方向的研究者与工程师提供风力发电机风扇叶片的语义分割数据集及配套Python训练代码可用于像素级识别叶片正常区域、磨损、裂缝与污渍等状况为风电运维提供决策支持。压缩包共2000个文件以1994个tif图像及对应标签图为主另有6个py脚本整体约810.93MB涵盖数据加载、预处理、数据集划分、模型训练与预测等完整环节。其中Unet.py实现了经典U-Net网络结构dataset.py定义数据加载方式pre_process.py负责图像预处理split_data.py完成训练验证集划分train.py与predict.py分别承担训练和推理任务形成可直接复用的语义分割流程。目前已有113人学习下载。数据集覆盖多种工作环境与光照条件标注经过筛选便于快速验证算法并缩短研发周期适合入门语义分割或开展风电叶片缺陷检测的读者参考使用。1. 风力发电机风扇语义分割从数据集到 Python 训练代码的落地路径拿到「风力发电机风扇语义分割数据集包含python训练代码」这个标题时我第一反应不是去找现成的压缩包而是先想清楚它到底解决什么问题。风电巡检场景里叶片和轮毂的缺陷检测、结冰识别、无人机航线规划都依赖一个前置能力把风扇本体从复杂背景里干净地抠出来。语义分割做的就是这个事——给每个像素打标签区分「风扇」和「非风扇」。这套数据集加训练代码的价值在于让你跳过最耗时的标注环节直接进入模型调参和部署验证。适合两类人一是做工业巡检算法、手头有风电图像但缺标注的工程师二是想拿一个真实工业场景练语义分割全流程的算法同学。下面我按「数据怎么组织、模型怎么选、代码怎么跑、坑在哪」的顺序把这条链路拆开讲。2. 风力发电机风扇语义分割数据集的结构与标注规范2.1 语义分割数据集长什么样目录、掩码与类别定义风力发电机风扇语义分割数据集常见做法是遵循 PASCAL VOC 或 Cityscapes 的组织习惯。核心就两块原图JPEG/PNG和对应的掩码图mask单通道 PNG。掩码里像素值直接代表类别比如 0 是背景1 是风扇叶片2 是轮毂3 是塔筒。有些数据集只做二分类风扇1背景0这种最省事也最适合刚上手语义分割的团队。目录结构我一般会整理成这样dataset/ ├── images/ │ ├── train/ │ │ ├── wind_0001.jpg │ │ └── ... │ └── val/ │ ├── wind_0101.jpg │ └── ... ├── masks/ │ ├── train/ │ │ ├── wind_0001.png │ │ └── ... │ └── val/ │ ├── wind_0101.png │ └── ... └── class_names.txtclass_names.txt里一行一个类别名顺序和掩码像素值对应。这个文件看着不起眼但后面做可视化调色板、算混淆矩阵都靠它别省。掩码的生成方式决定了数据质量。如果是人工标注边缘往往有 1~2 像素的抖动如果是用传统图像处理比如阈值形态学半自动生成叶片和天空对比度低的地方容易漏标。拿到数据集后我习惯先抽 20 张掩码叠在原图上肉眼过一遍重点看叶片尖端和轮毂连接处——这两个位置是语义分割模型最容易翻车的地方。2.2 标注质量自检三个必须跑的统计脚本数据集到手别急着训练先跑统计。下面这段 Python 代码做三件事统计每类像素占比、检查图像与掩码是否一一对应、找出掩码里出现但 class_names 没定义的异常像素值。import os import numpy as np from PIL import Image from collections import Counter IMG_DIR dataset/images/train MASK_DIR dataset/masks/train CLASS_NAMES [background, blade, hub, tower] def check_dataset(img_dir, mask_dir, class_names): img_files sorted(os.listdir(img_dir)) mask_files sorted(os.listdir(mask_dir)) # 检查文件名是否一一对应去掉扩展名比较 img_stems {os.path.splitext(f)[0] for f in img_files} mask_stems {os.path.splitext(f)[0] for f in mask_files} only_img img_stems - mask_stems only_mask mask_stems - img_stems if only_img: print(f[警告] 有图无掩码: {list(only_img)[:5]}) if only_mask: print(f[警告] 有掩码无图: {list(only_mask)[:5]}) pixel_counter Counter() for f in mask_files: mask np.array(Image.open(os.path.join(mask_dir, f))) pixel_counter.update(mask.flatten().tolist()) total sum(pixel_counter.values()) print(类别像素占比:) for idx, name in enumerate(class_names): cnt pixel_counter.get(idx, 0) print(f {idx} {name}: {cnt/total*100:.2f}%) # 找出未定义的像素值 defined set(range(len(class_names))) unknown set(pixel_counter.keys()) - defined if unknown: print(f[警告] 掩码中存在未定义像素值: {unknown}) check_dataset(IMG_DIR, MASK_DIR, CLASS_NAMES)逻辑说明img_stems和mask_stems用集合做差能快速定位缺失配对。pixel_counter统计所有掩码的像素分布如果某一类占比低于 1%说明样本极少训练时要么过采样要么在损失函数里给类别权重。unknown检查是血泪经验——有些数据集掩码用 255 表示前景但 class_names 只写了 0 和 1直接训练会导致标签越界模型输出全是乱码。参数方面CLASS_NAMES的顺序必须和掩码像素值严格对应改顺序等于改标签含义。如果数据集是二分类把列表改成[background, fan]即可。这个脚本跑完你对数据集的底细就有数了。3. 用 Python 训练风力发电机风扇语义分割模型从 DataLoader 到推理3.1 模型选型U-Net、DeepLabV3 还是 SegFormer风力发电机风扇的语义分割目标形状相对固定——叶片是细长条轮毂是圆形塔筒是竖直矩形。这种场景下模型选型不用追最贵的。我一般按数据量分三档数据量推荐模型理由 500 张U-Net ResNet34 backbone参数量小小样本下不容易过拟合500~2000 张DeepLabV3 MobileNetV3空洞卷积扩大感受野适合叶片这种细长目标 2000 张SegFormer (MiT-B2)Transformer 全局建模边缘更干净如果数据集里风扇只占图像一小部分比如无人机航拍风扇在画面角落DeepLabV3 的多尺度 ASPP 模块比 U-Net 更稳。反过来如果图像里风扇占了大半画面U-Net 的跳跃连接足够用训练还快。这里给一个基于segmentation_models_pytorch的 DeepLabV3 训练代码这个库封装了主流分割模型改 backbone 只要换一个字符串。import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import segmentation_models_pytorch as smp from PIL import Image import numpy as np import os import albumentations as A from albumentations.pytorch import ToTensorV2 class WindFanDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform self.images sorted(os.listdir(img_dir)) def __len__(self): return len(self.images) def __getitem__(self, idx): img_name self.images[idx] stem os.path.splitext(img_name)[0] img np.array(Image.open(os.path.join(self.img_dir, img_name)).convert(RGB)) mask np.array(Image.open(os.path.join(self.mask_dir, stem .png))) if self.transform: augmented self.transform(imageimg, maskmask) img augmented[image] mask augmented[mask] return img, mask.long() # 训练集增强翻转、旋转、颜色抖动 train_transform A.Compose([ A.Resize(512, 512), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.3), A.RandomRotate90(p0.5), A.ColorJitter(brightness0.2, contrast0.2, p0.3), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) val_transform A.Compose([ A.Resize(512, 512), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) train_ds WindFanDataset(dataset/images/train, dataset/masks/train, train_transform) val_ds WindFanDataset(dataset/images/val, dataset/masks/val, val_transform) train_loader DataLoader(train_ds, batch_size8, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size8, shuffleFalse, num_workers4) # 模型DeepLabV34 类输出 model smp.DeepLabV3Plus( encoder_namemobilenet_v3_large, encoder_weightsimagenet, in_channels3, classes4, ) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) # 损失函数交叉熵 Dice缓解类别不平衡 ce_loss nn.CrossEntropyLoss() dice_loss smp.losses.DiceLoss(modemulticlass) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50) def train_one_epoch(epoch): model.train() total_loss 0 for imgs, masks in train_loader: imgs, masks imgs.to(device), masks.to(device) optimizer.zero_grad() outputs model(imgs) loss ce_loss(outputs, masks) dice_loss(outputs, masks) loss.backward() optimizer.step() total_loss loss.item() scheduler.step() return total_loss / len(train_loader) for epoch in range(50): loss train_one_epoch(epoch) print(fEpoch {epoch1}, Loss: {loss:.4f})逻辑说明WindFanDataset把图像和掩码按文件名配对albumentations做同步增强——图像翻转时掩码跟着翻转这是分割任务和分类任务最大的区别。smp.DeepLabV3Plus的encoder_weightsimagenet表示用预训练权重小数据集上这步能省 30% 左右的收敛时间。损失函数用交叉熵加 Dice是因为风扇叶片像素占比通常远小于背景纯交叉熵会让模型倾向于全预测背景Dice 能拉回前景的梯度。参数方面batch_size8是 8GB 显存下的保守值显存够可以加到 16。lr1e-4配合 AdamW 和余弦退火是分割任务比较稳的组合。Resize(512, 512)是输入尺寸如果叶片在图像里很细可以提到 768但显存和训练时间会涨。3.2 训练过程监控IoU、Dice 与可视化验证训练 loss 下降不代表模型能用。语义分割必须看 IoU交并比和 Dice 系数而且最好每几个 epoch 存一张预测掩码叠在原图上的可视化图。下面这段代码算验证集 IoU 并保存对比图。import torch.nn.functional as F import matplotlib.pyplot as plt def evaluate(model, loader, device, num_classes4): model.eval() intersection torch.zeros(num_classes) union torch.zeros(num_classes) with torch.no_grad(): for imgs, masks in loader: imgs, masks imgs.to(device), masks.to(device) outputs model(imgs) preds torch.argmax(outputs, dim1) for cls in range(num_classes): pred_cls (preds cls) mask_cls (masks cls) intersection[cls] (pred_cls mask_cls).sum().item() union[cls] (pred_cls | mask_cls).sum().item() iou intersection / (union 1e-6) return iou def visualize_prediction(model, dataset, idx, device, save_pathpred_vis.png): model.eval() img, mask dataset[idx] with torch.no_grad(): output model(img.unsqueeze(0).to(device)) pred torch.argmax(output, dim1).squeeze(0).cpu().numpy() img_np img.permute(1, 2, 0).cpu().numpy() img_np (img_np - img_np.min()) / (img_np.max() - img_np.min()) fig, axes plt.subplots(1, 3, figsize(12, 4)) axes[0].imshow(img_np) axes[0].set_title(Image) axes[1].imshow(mask.cpu().numpy()) axes[1].set_title(GT Mask) axes[2].imshow(pred) axes[2].set_title(Prediction) for ax in axes: ax.axis(off) plt.savefig(save_path, dpi150, bbox_inchestight) plt.close() iou evaluate(model, val_loader, device) print(Per-class IoU:, iou) visualize_prediction(model, val_ds, idx0, devicedevice)逻辑说明intersection和union按类别累加最后逐类算 IoU。如果某一类 IoU 明显低比如轮毂只有 0.3说明该类样本太少或标注质量差需要针对性补数据。visualize_prediction把原图、真值掩码、预测掩码并排画出来叶片边缘是否平滑、有没有断裂一眼就能看出来。参数方面num_classes要和模型输出通道数一致。idx0可以换成任意验证集索引建议每次评估换不同索引避免只看一张图产生错觉。4. 风力发电机风扇语义分割的避坑与排查清单4.1 掩码像素值越界导致 loss 变 NaN现象训练第一个 epoch loss 就是 NaN或者模型输出全是同一个类别。原因掩码里出现了 class_names 未定义的像素值比如背景是 0、风扇是 255但模型只输出 2 类。CrossEntropyLoss遇到大于等于类别数的标签会直接报错或产生 NaN。解决跑 2.2 节的统计脚本把未知像素值映射到合法范围。如果原数据集用 255 表示前景训练前统一转成 1。mask np.array(Image.open(mask_path)) mask (mask 0).astype(np.uint8) # 255 - 1, 0 - 04.2 图像与掩码增强不同步导致标签错位现象训练 loss 能降但验证 IoU 极低可视化预测图和真值完全对不上。原因用了只对图像做增强的 pipeline掩码没跟着翻转或旋转。分割任务里图像和掩码必须共享同一套几何变换。解决用albumentations的Compose同时传入image和mask不要用torchvision.transforms分别处理。检查增强后掩码的类别分布是否和原图一致。4.3 类别极度不平衡导致模型只预测背景现象风扇 IoU 接近 0但背景 IoU 接近 1整体 accuracy 看起来很高。原因风扇像素占比可能只有 5%交叉熵损失被背景主导。解决损失函数加 Dice 或 Focal Loss或者在CrossEntropyLoss里传weight参数给前景类更高权重。权重可以按类别频率的倒数来设。class_counts np.array([800000, 50000, 20000, 30000]) # 统计得到的像素数 weights 1.0 / (class_counts 1e-6) weights weights / weights.sum() * len(weights) ce_loss nn.CrossEntropyLoss(weighttorch.tensor(weights, dtypetorch.float32).to(device))4.4 输入尺寸与显存不匹配导致训练中断现象训练到一半报CUDA out of memory或者 batch size 调到 2 还是爆显存。原因DeepLabV3 在 512×512 输入下显存占用比分类模型高很多尤其是 ASPP 模块。解决先把输入降到 384×384 跑通再逐步往上加。或者用梯度累积模拟大 batchbatch_size4累积 4 次再更新一次参数。混合精度训练也能省 30% 左右显存。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): outputs model(imgs) loss ce_loss(outputs, masks) dice_loss(outputs, masks) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()4.5 验证集指标虚高但实际部署效果差现象验证集 IoU 0.85但拿无人机实拍图推理风扇边缘一塌糊涂。原因验证集和训练集来自同一批数据分布太接近。实拍图的光照、角度、背景差异没被覆盖。解决划验证集时按拍摄批次或场景分不要随机抽。如果数据集里全是晴天图像自己补几张阴天、逆光的图做测试。另外推理时把输入尺寸调到和训练一致不要用原图直接跑。5. 把训练好的风扇分割模型推到推理端三个实用技巧模型训练完只是半成品真正落地还要过推理这一关。第一个技巧是滑动窗口推理。风力发电机图像往往分辨率很高无人机原图可能 4000×3000直接缩到 512 会丢失叶片细节。我一般把大图切成 512×512 的块块之间重叠 64 像素每块单独推理后再拼回去重叠区域取概率平均。这样叶片尖端不会被切丢。def sliding_window_inference(model, image, window512, stride448, num_classes4): model.eval() h, w image.shape[:2] prob_map np.zeros((num_classes, h, w), dtypenp.float32) count_map np.zeros((h, w), dtypenp.float32) for y in range(0, h, stride): for x in range(0, w, stride): y2 min(y window, h) x2 min(x window, w) y1 max(0, y2 - window) x1 max(0, x2 - window) patch image[y1:y2, x1:x2] # 预处理归一化、转 tensor、加 batch 维度 patch_tensor preprocess(patch).unsqueeze(0).to(device) with torch.no_grad(): out model(patch_tensor) prob torch.softmax(out, dim1).squeeze(0).cpu().numpy() prob_map[:, y1:y2, x1:x2] prob count_map[y1:y2, x1:x2] 1 prob_map / np.maximum(count_map, 1) return np.argmax(prob_map, axis0)window512和stride448意味着重叠 64 像素重叠越多边缘越平滑但推理时间线性增长。实际用的时候先拿一张图测看拼接缝明不明显再决定 stride。第二个技巧是掩码后处理。语义分割输出常有零星噪点用连通域分析去掉面积小于阈值的区域再对叶片做一次形态学闭运算边缘会干净很多。import cv2 from scipy import ndimage def postprocess_mask(mask, min_area200): # 去掉小连通域 labeled, num ndimage.label(mask) for i in range(1, num 1): if (labeled i).sum() min_area: mask[labeled i] 0 # 闭运算连接断裂的叶片 kernel cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) mask cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) return maskmin_area200是经验值图像分辨率高就调大。闭运算核大小 5×5 适合 512 输入如果叶片很细核再小一点否则会把背景也连进来。第三个技巧是用 TorchScript 或 ONNX 导出把推理速度提上来。PyTorch 原生推理在 Python 里跑每次都有解释开销。导出成 ONNX 后用onnxruntime跑CPU 上也能快 2~3 倍。dummy_input torch.randn(1, 3, 512, 512).to(device) torch.onnx.export( model, dummy_input, wind_fan_deeplabv3.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version11, )导出后记得用onnxruntime跑一遍验证输出和 PyTorch 一致数值误差在 1e-4 以内算正常。如果部署在边缘设备上还可以进一步做 INT8 量化但量化后一定要重新评估 IoU有些模型量化后边缘会变毛糙。最后说个我自己的习惯每次训练完我会把验证集里 IoU 最低的 10 张图单独存一个文件夹逐张看预测结果。这 10 张图往往暴露了数据集中最棘手的问题——可能是逆光、可能是叶片和背景颜色接近、可能是标注本身就有误。改完这些再重训比盲目加数据有效得多。希望帮到你。本文还有配套的精品资源点击获取
返回列表