電機葉片語義分割實戰(zhàn):U-Net數(shù)據(jù)集與訓練全流程)
簡介本資源為風力發(fā)電機風扇葉片語義分割數(shù)據(jù)集面向從事計算機視覺與智能風電運維的研究者、工程師及學生用于訓練和驗證像素級葉片狀態(tài)識別模型可區(qū)分正常區(qū)域、磨損、裂縫與污漬等狀況。壓縮包共約2000個文件以1994個tif圖像及對應標簽圖為主另含6個Python腳本整體約810.93MB涵蓋數(shù)據(jù)加載、圖像預處理、數(shù)據(jù)集劃分、U-Net網(wǎng)絡(luò)實現(xiàn)、模型訓練與預測等完整流程目錄結(jié)構(gòu)清晰便于按模塊檢索與二次開發(fā)。目前已有113人學習下載。數(shù)據(jù)集覆蓋多種工作環(huán)境與光照條件配合開箱即用的訓練代碼讀者可快速搭建語義分割實驗、復現(xiàn)U-Net基線并遷移到自有數(shù)據(jù)為風電葉片智能監(jiān)測提供可落地的算法框架與排錯參考。1. 風力發(fā)電機葉片語義分割數(shù)據(jù)集從一堆 tif 到能跑通的 U-Net拿到這個資源包的時候我第一反應是這數(shù)據(jù)集有點意思。目錄里躺著train.py、Unet.py、split_data.py、predict.py、pre_process.py、dataset.py六個腳本外加一個DATASET文件夾里面是成對的葉片圖像和標簽圖格式是.tif。這不是那種只給你一堆圖片讓你自己猜標注格式的半成品數(shù)據(jù)集而是把預處理、劃分、訓練、預測整條鏈路都鋪好了。它解決的核心問題很明確你想驗證一個語義分割模型在工業(yè)葉片缺陷場景下到底行不行不用從零標注、不用自己搭 U-Net改改路徑就能跑。適合誰做風電運維智能巡檢的算法工程師、拿語義分割練手的學生、以及想快速對比 U-Net 和其他分割網(wǎng)絡(luò)在真實工業(yè)數(shù)據(jù)上表現(xiàn)的從業(yè)者。下面我按數(shù)據(jù)長什么樣 → 代碼怎么串 → 坑在哪 → 怎么改的順序拆一遍。2. 數(shù)據(jù)集結(jié)構(gòu)與預處理鏈路tif 格式、標簽對齊與 split_data 的劃分邏輯2.1 為什么是 tif 而不是 jpg資源里的圖像文件后綴是.tif比如TCGA_DU_6408_19860521_25.tif這種命名。tif 在工業(yè)圖像里常見原因是它支持無損壓縮和多通道存儲葉片表面的細微裂紋、污漬在 jpg 壓縮下容易產(chǎn)生塊效應標注邊界會糊。語義分割對邊界敏感尤其是裂縫這種細長目標壓縮偽影會直接拉低 IoU。所以拿到 tif 不要急著轉(zhuǎn) jpg先確認你的數(shù)據(jù)加載鏈路能不能直接讀。常見做法是用Pillow或OpenCV讀但兩者對多通道 tif 的處理有差異。我一般統(tǒng)一用cv2.imread(path, cv2.IMREAD_UNCHANGED)保留原始位深避免自動轉(zhuǎn) 8 位丟信息。如果顯存吃緊再在pre_process.py里做 resize而不是在讀取階段就降質(zhì)。2.2 pre_process.py 里該盯住的三個參數(shù)pre_process.py通常干三件事統(tǒng)一尺寸、歸一化、可選的數(shù)據(jù)增強。這個腳本沒有在正文里給出具體實現(xiàn)但按語義分割的常規(guī)做法我會重點檢查這幾個點輸入尺寸U-Net 對輸入尺寸有下采樣倍數(shù)要求一般是 16 或 32 的整數(shù)倍。如果原圖是 512×512 那沒問題如果是任意尺寸resize 到 256×256 或 512×512 最穩(wěn)。歸一化方式用 ImageNet 均值方差還是簡單除以 255取決于你是否加載預訓練權(quán)重。從零訓練就用/255.0加載預訓練就對齊mean[0.485,0.456,0.406]。標簽處理語義分割標簽必須是單通道類別索引圖不是 RGB 彩圖。如果DATASET里的標簽是彩色掩膜需要在dataset.py里做顏色到類別的映射否則CrossEntropyLoss會直接報維度錯誤。import cv2 import numpy as np def preprocess_image(img_path, target_size(512, 512)): # IMREAD_UNCHANGED 保留原始位深避免 tif 被截斷 img cv2.imread(img_path, cv2.IMREAD_UNCHANGED) if img is None: raise FileNotFoundError(f讀不到圖: {img_path}) # 統(tǒng)一到 3 通道灰度圖也能進 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。標簽圖 resize 用線性插值會把類別 1 和類別 2 之間插出 1.5訓練時 loss 直接崩。這個坑我見過不止一次現(xiàn)象是 loss 一開始就是 NaN查半天以為是學習率其實是標簽被插值污染了。2.3 split_data.py 的劃分比例與隨機種子split_data.py負責把DATASET切成訓練集、驗證集、測試集。工業(yè)數(shù)據(jù)集樣本量通常不大劃分比例我一般用 7:2:1 或 8:1:1。重點不是比例是隨機種子固定和按圖像對劃分。如果腳本是按單張圖隨機抽可能出現(xiàn)訓練集里有原圖、驗證集里有對應標簽的情況指標虛高。正確做法是以圖像-標簽對為單位劃分。另外如果同一葉片拍了多張不同角度的圖最好按葉片 ID 分組劃分避免同一葉片同時出現(xiàn)在訓練和驗證集里這叫數(shù)據(jù)泄漏是語義分割里最隱蔽的翻車點之一。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)) # 只保留有對應標簽的圖防止訓練時找不到 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固定后每次運行結(jié)果一致方便復現(xiàn)實驗。pairs的過濾邏輯很重要如果DATASET里存在沒有標簽的圖不過濾就會在訓練時拋FileNotFoundError。3. U-Net 網(wǎng)絡(luò)結(jié)構(gòu)與訓練腳本Unet.py 的通道設(shè)計、train.py 的損失與學習率3.1 Unet.py 的編碼器-解碼器通道數(shù)怎么定Unet.py實現(xiàn)的是經(jīng)典 U-Net編碼器逐層下采樣提特征解碼器逐層上采樣恢復分辨率中間用跳躍連接把淺層細節(jié)拼回來。語義分割里 U-Net 的優(yōu)勢就在跳躍連接葉片裂縫這種細目標沒有淺層特征補充解碼器根本恢復不出邊界。通道數(shù)常見配置是[64, 128, 256, 512, 1024]但這是針對 512×512 輸入的。如果你的圖 resize 到 256×256可以砍到[32, 64, 128, 256]顯存占用能降一半以上。改通道數(shù)的時候注意編碼器每下采樣一次通道翻倍解碼器每上采樣一次通道減半跳躍連接拼接時通道要對齊否則torch.cat會報維度不匹配。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ù)要對應 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要按你的標簽類別數(shù)改。如果標簽里只有背景和葉片兩類就是 2如果有背景、正常葉片、磨損、裂縫四類就是 4。這個值必須和dataset.py里標簽映射的類別數(shù)一致不一致的話訓練不報錯但預測結(jié)果全黑因為模型輸出的通道和標簽索引對不上。3.2 train.py 的損失函數(shù)與學習率策略train.py是整條鏈路的核心。語義分割常用CrossEntropyLoss如果類別極度不平衡比如裂縫像素只占 1%換成DiceLoss或CrossEntropyLoss DiceLoss組合。學習率我一般從1e-3起步配CosineAnnealingLR或ReduceLROnPlateau前者平滑下降后者根據(jù)驗證集指標動態(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() # 驗證階段只算 loss指標另算 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 啟動階段這是 Windows 下 PyTorch 的經(jīng)典坑。T_max50要和總 epoch 數(shù)一致不然余弦退火曲線會提前走完。3.3 dataset.py 的標簽映射與增強邊界dataset.py負責把圖像和標簽配對讀進來轉(zhuǎn)成 tensor。這里最容易出問題的是標簽映射如果標簽圖里像素值是 0 和 255而你的num_classes2模型期望的是 0 和 1直接送進去 loss 會算錯。需要在__getitem__里做mask (mask 0).long()或按顏色映射表轉(zhuǎn)換。數(shù)據(jù)增強方面圖像和標簽必須同步變換。翻轉(zhuǎn)、旋轉(zhuǎn)可以同步做但顏色抖動只能作用于圖像不能碰標簽。如果dataset.py里用了albumentations記得把additional_targets{mask: mask}加上否則標簽不會跟著變換訓練出來的模型邊界全是錯的。4. 避坑與排查tif 讀取、顯存溢出、標簽錯位、指標虛高4.1 現(xiàn)象訓練一開始 loss 就是 NaN原因標簽圖 resize 用了線性插值類別索引被插成小數(shù)CrossEntropyLoss要求標簽是int64小數(shù)被截斷后類別錯亂。或者圖像歸一化時除了 0某些像素全黑導致除零。解決標簽 resize 強制用INTER_NEAREST歸一化前檢查圖像像素范圍加np.clip兜底。在dataset.py里打印一次 mask 的unique()值確認只有預期類別。4.2 現(xiàn)象CUDA out of memory但 batch_size 已經(jīng)調(diào)到 1原因U-Net 的跳躍連接在拼接時顯存占用翻倍512×512 輸入下 base_ch64 的模型batch_size1 也可能吃滿 8G。另外num_workers過大導致每個 worker 都復制一份數(shù)據(jù)到顯存。解決把輸入 resize 到 256×256或把base_ch從 64 降到 32。訓練時用torch.cuda.empty_cache()清理緩存num_workers設(shè) 2 以內(nèi)。如果還不夠用混合精度訓練torch.cuda.amp顯存能省 30% 左右。4.3 現(xiàn)象驗證集 IoU 很高但 predict.py 跑出來全是背景原因類別不平衡。裂縫像素占比極低模型學會全預測背景就能拿到高準確率但 IoU 對少數(shù)類不敏感。或者predict.py里的閾值設(shè)錯了argmax 之后沒有做后處理。解決換DiceLoss或加類別權(quán)重weighttorch.tensor([1.0, 10.0])。predict.py里輸出 argmax 后檢查每個類別的像素占比如果少數(shù)類占比為 0說明模型沒學到需要回頭查標簽映射和損失函數(shù)。4.4 現(xiàn)象split_data.py 跑完訓練集和驗證集有同名文件原因腳本按文件名排序后直接切片如果DATASET里圖像和標簽分兩個文件夾且文件名不完全對應切片后可能錯位。解決劃分前先做一次配對校驗只保留圖像和標簽都存在的文件對。劃分后打印三集的文件名列表肉眼掃一遍有沒有重疊。這個檢查花不了一分鐘但能省掉后面幾小時的無效訓練。4.5 現(xiàn)象Windows 下 DataLoader 卡死不動原因num_workers 0時Windows 的進程啟動方式與 Linux 不同PyTorch 的 DataLoader 在 Windows 上容易死鎖。解決把num_workers設(shè)為 0或者把訓練代碼包在if __name__ __main__:里。后者是標準做法但很多人寫腳本時忘了加導致多進程反復啟動主模塊。5. 進階用法從 predict.py 到指標驗證以及我改 U-Net 的一個習慣predict.py不只是拿來看效果的它可以改成批量推理腳本輸出每張圖的預測掩膜和置信度再和標簽對比算 IoU、Dice、Precision、Recall。我一般會在predict.py里加一段指標計算而不是另寫腳本因為預測和評估用的是同一套預處理邏輯分開寫容易不一致。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) # 實際用 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對每個類別單獨算再平均比整體像素準確率更能反映分割質(zhì)量。如果某個類別 union 為 0說明測試集里沒有這個類跳過而不是算 0否則指標會被拉低。我改 U-Net 的一個習慣在Unet.py的out層之前加一個Dropout2d(0.1)。工業(yè)數(shù)據(jù)集樣本少模型容易過擬合加輕量 dropout 后驗證集 loss 通常更穩(wěn)。這個改動不影響推理因為model.eval()時 dropout 自動關(guān)閉。另外如果DATASET里的葉片圖像背景復雜可以在dataset.py里加隨機裁剪讓模型多見局部區(qū)域?qū)α芽p這種小目標提升明顯。從那以后我每次拿到新的分割數(shù)據(jù)集都強制先跑一遍split_data.py的配對校驗和dataset.py的 mask unique 檢查再開始訓練。這兩個檢查加起來不到五分鐘但能擋掉后面大部分的玄學問題。希望幫到你。本文還有配套的精品資源點擊獲取