力發(fā)電機(jī)葉片語義分割實(shí)戰(zhàn):U-Net數(shù)據(jù)集與訓(xùn)練全流程)
簡介本資源為風(fēng)力發(fā)電機(jī)風(fēng)扇葉片語義分割數(shù)據(jù)集面向從事計(jì)算機(jī)視覺與智能風(fēng)電運(yùn)維的研究者、工程師及學(xué)生用于訓(xùn)練和驗(yàn)證像素級(jí)葉片狀態(tài)識(shí)別模型可區(qū)分正常區(qū)域、磨損、裂縫與污漬等狀況。壓縮包共約2000個(gè)文件以1994個(gè)tif圖像及對應(yīng)標(biāo)簽圖為主另含6個(gè)Python腳本整體約810.93MB涵蓋數(shù)據(jù)加載、圖像預(yù)處理、數(shù)據(jù)集劃分、U-Net網(wǎng)絡(luò)實(shí)現(xiàn)、模型訓(xùn)練與預(yù)測等完整流程目錄結(jié)構(gòu)清晰便于按模塊檢索與二次開發(fā)。目前已有113人學(xué)習(xí)下載。數(shù)據(jù)集覆蓋多種工作環(huán)境與光照條件配合開箱即用的訓(xùn)練代碼讀者可快速搭建語義分割實(shí)驗(yàn)、復(fù)現(xiàn)U-Net基線并遷移到自有數(shù)據(jù)為風(fēng)電葉片智能監(jiān)測提供可落地的算法框架與排錯(cuò)參考。1. 風(fēng)力發(fā)電機(jī)葉片語義分割數(shù)據(jù)集從一堆 tif 到能跑通的 U-Net拿到這個(gè)資源包的時(shí)候我第一反應(yīng)是這數(shù)據(jù)集有點(diǎn)意思。目錄里躺著train.py、Unet.py、split_data.py、predict.py、pre_process.py、dataset.py六個(gè)腳本外加一個(gè)DATASET文件夾里面是成對的葉片圖像和標(biāo)簽圖格式是.tif。這不是那種只給你一堆圖片讓你自己猜標(biāo)注格式的半成品數(shù)據(jù)集而是把預(yù)處理、劃分、訓(xùn)練、預(yù)測整條鏈路都鋪好了。它解決的核心問題很明確你想驗(yàn)證一個(gè)語義分割模型在工業(yè)葉片缺陷場景下到底行不行不用從零標(biāo)注、不用自己搭 U-Net改改路徑就能跑。適合誰做風(fēng)電運(yùn)維智能巡檢的算法工程師、拿語義分割練手的學(xué)生、以及想快速對比 U-Net 和其他分割網(wǎng)絡(luò)在真實(shí)工業(yè)數(shù)據(jù)上表現(xiàn)的從業(yè)者。下面我按數(shù)據(jù)長什么樣 → 代碼怎么串 → 坑在哪 → 怎么改的順序拆一遍。2. 數(shù)據(jù)集結(jié)構(gòu)與預(yù)處理鏈路tif 格式、標(biāo)簽對齊與 split_data 的劃分邏輯2.1 為什么是 tif 而不是 jpg資源里的圖像文件后綴是.tif比如TCGA_DU_6408_19860521_25.tif這種命名。tif 在工業(yè)圖像里常見原因是它支持無損壓縮和多通道存儲(chǔ)葉片表面的細(xì)微裂紋、污漬在 jpg 壓縮下容易產(chǎn)生塊效應(yīng)標(biāo)注邊界會(huì)糊。語義分割對邊界敏感尤其是裂縫這種細(xì)長目標(biāo)壓縮偽影會(huì)直接拉低 IoU。所以拿到 tif 不要急著轉(zhuǎn) jpg先確認(rèn)你的數(shù)據(jù)加載鏈路能不能直接讀。常見做法是用Pillow或OpenCV讀但兩者對多通道 tif 的處理有差異。我一般統(tǒng)一用cv2.imread(path, cv2.IMREAD_UNCHANGED)保留原始位深避免自動(dòng)轉(zhuǎn) 8 位丟信息。如果顯存吃緊再在pre_process.py里做 resize而不是在讀取階段就降質(zhì)。2.2 pre_process.py 里該盯住的三個(gè)參數(shù)pre_process.py通常干三件事統(tǒng)一尺寸、歸一化、可選的數(shù)據(jù)增強(qiáng)。這個(gè)腳本沒有在正文里給出具體實(shí)現(xiàn)但按語義分割的常規(guī)做法我會(huì)重點(diǎn)檢查這幾個(gè)點(diǎn)輸入尺寸U-Net 對輸入尺寸有下采樣倍數(shù)要求一般是 16 或 32 的整數(shù)倍。如果原圖是 512×512 那沒問題如果是任意尺寸resize 到 256×256 或 512×512 最穩(wěn)。歸一化方式用 ImageNet 均值方差還是簡單除以 255取決于你是否加載預(yù)訓(xùn)練權(quán)重。從零訓(xùn)練就用/255.0加載預(yù)訓(xùn)練就對齊mean[0.485,0.456,0.406]。標(biāo)簽處理語義分割標(biāo)簽必須是單通道類別索引圖不是 RGB 彩圖。如果DATASET里的標(biāo)簽是彩色掩膜需要在dataset.py里做顏色到類別的映射否則CrossEntropyLoss會(huì)直接報(bào)維度錯(cuò)誤。import cv2 import numpy as np def preprocess_image(img_path, target_size(512, 512)): # IMREAD_UNCHANGED 保留原始位深避免 tif 被截?cái)?img cv2.imread(img_path, cv2.IMREAD_UNCHANGED) if img is None: raise FileNotFoundError(f讀不到圖: {img_path}) # 統(tǒng)一到 3 通道灰度圖也能進(jìn) U-Net if len(img.shape) 2: img cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) img cv2.resize(img, target_size, interpolationcv2.INTER_LINEAR) img img.astype(np.float32) / 255.0 return img def preprocess_mask(mask_path, target_size(512, 512)): mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 最近鄰插值防止類別索引被線性插值搞出小數(shù) mask cv2.resize(mask, target_size, interpolationcv2.INTER_NEAREST) return mask.astype(np.int64)這段代碼的關(guān)鍵在INTER_NEAREST。標(biāo)簽圖 resize 用線性插值會(huì)把類別 1 和類別 2 之間插出 1.5訓(xùn)練時(shí) loss 直接崩。這個(gè)坑我見過不止一次現(xiàn)象是 loss 一開始就是 NaN查半天以為是學(xué)習(xí)率其實(shí)是標(biāo)簽被插值污染了。2.3 split_data.py 的劃分比例與隨機(jī)種子split_data.py負(fù)責(zé)把DATASET切成訓(xùn)練集、驗(yàn)證集、測試集。工業(yè)數(shù)據(jù)集樣本量通常不大劃分比例我一般用 7:2:1 或 8:1:1。重點(diǎn)不是比例是隨機(jī)種子固定和按圖像對劃分。如果腳本是按單張圖隨機(jī)抽可能出現(xiàn)訓(xùn)練集里有原圖、驗(yàn)證集里有對應(yīng)標(biāo)簽的情況指標(biāo)虛高。正確做法是以圖像-標(biāo)簽對為單位劃分。另外如果同一葉片拍了多張不同角度的圖最好按葉片 ID 分組劃分避免同一葉片同時(shí)出現(xiàn)在訓(xùn)練和驗(yàn)證集里這叫數(shù)據(jù)泄漏是語義分割里最隱蔽的翻車點(diǎn)之一。import os import random import shutil def split_dataset(img_dir, mask_dir, out_dir, ratios(0.7, 0.2, 0.1), seed42): random.seed(seed) # 固定種子保證每次劃分一致 files sorted(os.listdir(img_dir)) # 只保留有對應(yīng)標(biāo)簽的圖防止訓(xùn)練時(shí)找不到 mask pairs [f for f in files if os.path.exists(os.path.join(mask_dir, f))] random.shuffle(pairs) n len(pairs) n_train int(n * ratios[0]) n_val int(n * ratios[1]) splits { train: pairs[:n_train], val: pairs[n_train:n_train n_val], test: pairs[n_train n_val:] } for split, items in splits.items(): for sub in [images, masks]: os.makedirs(os.path.join(out_dir, split, sub), exist_okTrue) for f in items: shutil.copy(os.path.join(img_dir, f), os.path.join(out_dir, split, images, f)) shutil.copy(os.path.join(mask_dir, f), os.path.join(out_dir, split, masks, f)) print(f劃分完成: train{len(splits[train])}, val{len(splits[val])}, test{len(splits[test])})參數(shù)說明ratios控制三集比例seed固定后每次運(yùn)行結(jié)果一致方便復(fù)現(xiàn)實(shí)驗(yàn)。pairs的過濾邏輯很重要如果DATASET里存在沒有標(biāo)簽的圖不過濾就會(huì)在訓(xùn)練時(shí)拋FileNotFoundError。3. U-Net 網(wǎng)絡(luò)結(jié)構(gòu)與訓(xùn)練腳本Unet.py 的通道設(shè)計(jì)、train.py 的損失與學(xué)習(xí)率3.1 Unet.py 的編碼器-解碼器通道數(shù)怎么定Unet.py實(shí)現(xiàn)的是經(jīng)典 U-Net編碼器逐層下采樣提特征解碼器逐層上采樣恢復(fù)分辨率中間用跳躍連接把淺層細(xì)節(jié)拼回來。語義分割里 U-Net 的優(yōu)勢就在跳躍連接葉片裂縫這種細(xì)目標(biāo)沒有淺層特征補(bǔ)充解碼器根本恢復(fù)不出邊界。通道數(shù)常見配置是[64, 128, 256, 512, 1024]但這是針對 512×512 輸入的。如果你的圖 resize 到 256×256可以砍到[32, 64, 128, 256]顯存占用能降一半以上。改通道數(shù)的時(shí)候注意編碼器每下采樣一次通道翻倍解碼器每上采樣一次通道減半跳躍連接拼接時(shí)通道要對齊否則torch.cat會(huì)報(bào)維度不匹配。import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net 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.net(x) class UNet(nn.Module): def __init__(self, in_ch3, num_classes2, base_ch64): super().__init__() # 編碼器每層通道翻倍 self.enc1 DoubleConv(in_ch, base_ch) self.enc2 DoubleConv(base_ch, base_ch * 2) self.enc3 DoubleConv(base_ch * 2, base_ch * 4) self.enc4 DoubleConv(base_ch * 4, base_ch * 8) self.pool nn.MaxPool2d(2) # 瓶頸層 self.bottleneck DoubleConv(base_ch * 8, base_ch * 16) # 解碼器上采樣后與跳躍連接拼接通道數(shù)要對應(yīng) self.up4 nn.ConvTranspose2d(base_ch * 16, base_ch * 8, 2, stride2) self.dec4 DoubleConv(base_ch * 16, base_ch * 8) self.up3 nn.ConvTranspose2d(base_ch * 8, base_ch * 4, 2, stride2) self.dec3 DoubleConv(base_ch * 8, base_ch * 4) self.up2 nn.ConvTranspose2d(base_ch * 4, base_ch * 2, 2, stride2) self.dec2 DoubleConv(base_ch * 4, base_ch * 2) self.up1 nn.ConvTranspose2d(base_ch * 2, base_ch, 2, stride2) self.dec1 DoubleConv(base_ch * 2, base_ch) self.out nn.Conv2d(base_ch, num_classes, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)num_classes要按你的標(biāo)簽類別數(shù)改。如果標(biāo)簽里只有背景和葉片兩類就是 2如果有背景、正常葉片、磨損、裂縫四類就是 4。這個(gè)值必須和dataset.py里標(biāo)簽映射的類別數(shù)一致不一致的話訓(xùn)練不報(bào)錯(cuò)但預(yù)測結(jié)果全黑因?yàn)槟P洼敵龅耐ǖ篮蜆?biāo)簽索引對不上。3.2 train.py 的損失函數(shù)與學(xué)習(xí)率策略train.py是整條鏈路的核心。語義分割常用CrossEntropyLoss如果類別極度不平衡比如裂縫像素只占 1%換成DiceLoss或CrossEntropyLoss DiceLoss組合。學(xué)習(xí)率我一般從1e-3起步配CosineAnnealingLR或ReduceLROnPlateau前者平滑下降后者根據(jù)驗(yàn)證集指標(biāo)動(dòng)態(tài)調(diào)。import torch from torch.utils.data import DataLoader from dataset import BladeDataset from Unet import UNet device torch.device(cuda if torch.cuda.is_available() else cpu) model UNet(in_ch3, num_classes2, base_ch64).to(device) criterion torch.nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50) train_ds BladeDataset(DATASET/train, augmentTrue) val_ds BladeDataset(DATASET/val, augmentFalse) train_loader DataLoader(train_ds, batch_size4, shuffleTrue, num_workers2) val_loader DataLoader(val_ds, batch_size4, shuffleFalse, num_workers2) for epoch in range(50): model.train() for imgs, masks in train_loader: imgs, masks imgs.to(device), masks.to(device) optimizer.zero_grad() logits model(imgs) loss criterion(logits, masks) loss.backward() optimizer.step() scheduler.step() # 驗(yàn)證階段只算 loss指標(biāo)另算 model.eval() with torch.no_grad(): for imgs, masks in val_loader: imgs, masks imgs.to(device), masks.to(device) _ model(imgs) print(fepoch {epoch} done, lr{scheduler.get_last_lr()[0]:.6f})batch_size4是 8G 顯存下的保守值顯存夠可以加到 8 或 16。num_workers在 Windows 上建議設(shè) 0否則容易卡在 DataLoader 啟動(dòng)階段這是 Windows 下 PyTorch 的經(jīng)典坑。T_max50要和總 epoch 數(shù)一致不然余弦退火曲線會(huì)提前走完。3.3 dataset.py 的標(biāo)簽映射與增強(qiáng)邊界dataset.py負(fù)責(zé)把圖像和標(biāo)簽配對讀進(jìn)來轉(zhuǎn)成 tensor。這里最容易出問題的是標(biāo)簽映射如果標(biāo)簽圖里像素值是 0 和 255而你的num_classes2模型期望的是 0 和 1直接送進(jìn)去 loss 會(huì)算錯(cuò)。需要在__getitem__里做mask (mask 0).long()或按顏色映射表轉(zhuǎn)換。數(shù)據(jù)增強(qiáng)方面圖像和標(biāo)簽必須同步變換。翻轉(zhuǎn)、旋轉(zhuǎn)可以同步做但顏色抖動(dòng)只能作用于圖像不能碰標(biāo)簽。如果dataset.py里用了albumentations記得把a(bǔ)dditional_targets{mask: mask}加上否則標(biāo)簽不會(huì)跟著變換訓(xùn)練出來的模型邊界全是錯(cuò)的。4. 避坑與排查tif 讀取、顯存溢出、標(biāo)簽錯(cuò)位、指標(biāo)虛高4.1 現(xiàn)象訓(xùn)練一開始 loss 就是 NaN原因標(biāo)簽圖 resize 用了線性插值類別索引被插成小數(shù)CrossEntropyLoss要求標(biāo)簽是int64小數(shù)被截?cái)嗪箢悇e錯(cuò)亂。或者圖像歸一化時(shí)除了 0某些像素全黑導(dǎo)致除零。解決標(biāo)簽 resize 強(qiáng)制用INTER_NEAREST歸一化前檢查圖像像素范圍加np.clip兜底。在dataset.py里打印一次 mask 的unique()值確認(rèn)只有預(yù)期類別。4.2 現(xiàn)象CUDA out of memory但 batch_size 已經(jīng)調(diào)到 1原因U-Net 的跳躍連接在拼接時(shí)顯存占用翻倍512×512 輸入下 base_ch64 的模型batch_size1 也可能吃滿 8G。另外num_workers過大導(dǎo)致每個(gè) worker 都復(fù)制一份數(shù)據(jù)到顯存。解決把輸入 resize 到 256×256或把base_ch從 64 降到 32。訓(xùn)練時(shí)用torch.cuda.empty_cache()清理緩存num_workers設(shè) 2 以內(nèi)。如果還不夠用混合精度訓(xùn)練torch.cuda.amp顯存能省 30% 左右。4.3 現(xiàn)象驗(yàn)證集 IoU 很高但 predict.py 跑出來全是背景原因類別不平衡。裂縫像素占比極低模型學(xué)會(huì)全預(yù)測背景就能拿到高準(zhǔn)確率但 IoU 對少數(shù)類不敏感?;蛘遬redict.py里的閾值設(shè)錯(cuò)了argmax 之后沒有做后處理。解決換DiceLoss或加類別權(quán)重weighttorch.tensor([1.0, 10.0])。predict.py里輸出 argmax 后檢查每個(gè)類別的像素占比如果少數(shù)類占比為 0說明模型沒學(xué)到需要回頭查標(biāo)簽映射和損失函數(shù)。4.4 現(xiàn)象split_data.py 跑完訓(xùn)練集和驗(yàn)證集有同名文件原因腳本按文件名排序后直接切片如果DATASET里圖像和標(biāo)簽分兩個(gè)文件夾且文件名不完全對應(yīng)切片后可能錯(cuò)位。解決劃分前先做一次配對校驗(yàn)只保留圖像和標(biāo)簽都存在的文件對。劃分后打印三集的文件名列表肉眼掃一遍有沒有重疊。這個(gè)檢查花不了一分鐘但能省掉后面幾小時(shí)的無效訓(xùn)練。4.5 現(xiàn)象Windows 下 DataLoader 卡死不動(dòng)原因num_workers 0時(shí)Windows 的進(jìn)程啟動(dòng)方式與 Linux 不同PyTorch 的 DataLoader 在 Windows 上容易死鎖。解決把num_workers設(shè)為 0或者把訓(xùn)練代碼包在if __name__ __main__:里。后者是標(biāo)準(zhǔn)做法但很多人寫腳本時(shí)忘了加導(dǎo)致多進(jìn)程反復(fù)啟動(dòng)主模塊。5. 進(jìn)階用法從 predict.py 到指標(biāo)驗(yàn)證以及我改 U-Net 的一個(gè)習(xí)慣predict.py不只是拿來看效果的它可以改成批量推理腳本輸出每張圖的預(yù)測掩膜和置信度再和標(biāo)簽對比算 IoU、Dice、Precision、Recall。我一般會(huì)在predict.py里加一段指標(biāo)計(jì)算而不是另寫腳本因?yàn)轭A(yù)測和評估用的是同一套預(yù)處理邏輯分開寫容易不一致。import torch import numpy as np from Unet import UNet from pre_process import preprocess_image def compute_iou(pred, target, num_classes2): ious [] for cls in range(num_classes): inter ((pred cls) (target cls)).sum() union ((pred cls) | (target cls)).sum() # 少數(shù)類可能 union 為 0跳過避免除零 if union 0: continue ious.append(inter / union) return np.mean(ious) if ious else 0.0 model UNet(in_ch3, num_classes2).cuda() model.load_state_dict(torch.load(best_unet.pth)) model.eval() img preprocess_image(DATASET/test/images/sample.tif) mask preprocess_image(DATASET/test/masks/sample.tif) # 實(shí)際用 mask 讀取邏輯 with torch.no_grad(): logits model(torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).cuda()) pred logits.argmax(dim1).squeeze().cpu().numpy() print(fIoU: {compute_iou(pred, mask):.4f})這段代碼里compute_iou對每個(gè)類別單獨(dú)算再平均比整體像素準(zhǔn)確率更能反映分割質(zhì)量。如果某個(gè)類別 union 為 0說明測試集里沒有這個(gè)類跳過而不是算 0否則指標(biāo)會(huì)被拉低。我改 U-Net 的一個(gè)習(xí)慣在Unet.py的out層之前加一個(gè)Dropout2d(0.1)。工業(yè)數(shù)據(jù)集樣本少模型容易過擬合加輕量 dropout 后驗(yàn)證集 loss 通常更穩(wěn)。這個(gè)改動(dòng)不影響推理因?yàn)閙odel.eval()時(shí) dropout 自動(dòng)關(guān)閉。另外如果DATASET里的葉片圖像背景復(fù)雜可以在dataset.py里加隨機(jī)裁剪讓模型多見局部區(qū)域?qū)α芽p這種小目標(biāo)提升明顯。從那以后我每次拿到新的分割數(shù)據(jù)集都強(qiáng)制先跑一遍split_data.py的配對校驗(yàn)和dataset.py的 mask unique 檢查再開始訓(xùn)練。這兩個(gè)檢查加起來不到五分鐘但能擋掉后面大部分的玄學(xué)問題。希望幫到你。本文還有配套的精品資源點(diǎn)擊獲取