簡介面向Python開發(fā)者的圖像分割實戰(zhàn)資料圍繞U-Net模型講解如何完成像素級分割任務(wù)適用于遙感影像、醫(yī)學(xué)圖像等場景。壓縮包共21個文件、約5.6MB內(nèi)附1個Python腳本平滑切片預(yù)測、多張jpg/png樣例圖、2個gif動態(tài)對比圖以及Markdown說明文檔便于對照代碼與效果圖理解模型推理過程。已有10157人學(xué)習(xí)下載。資料內(nèi)容涵蓋U-Net收縮/擴(kuò)展路徑原理、數(shù)據(jù)掩模準(zhǔn)備、交叉熵與Dice損失選擇、優(yōu)化器配置等關(guān)鍵知識點(diǎn)代碼庫Smoothly-Blend-Image-Patches-master實現(xiàn)了圖像分塊平滑融合可用于消除切片邊界偽影提升分割效果。通過閱讀README并運(yùn)行腳本讀者能掌握從模型構(gòu)建、訓(xùn)練到預(yù)測的完整流程還可借助樣例圖像直觀了解U-Net在衛(wèi)星圖像分割中的實際表現(xiàn)適合希望以Python快速落地圖像分割項目的中級開發(fā)者。1. UNet 圖像分割把像素分類這件事講透圖像分割和分類最大的區(qū)別在于分類只回答“這張圖里有什么”而分割要回答“這個物體在哪、邊界在哪里、和旁邊的物體怎么分開”。UNet 之所以成為分割任務(wù)里的常青樹是因為它在編碼器-解碼器結(jié)構(gòu)上加了跳躍連接用很小的參數(shù)量同時保住了全局語義和局部細(xì)節(jié)——我做過不少分割項目從路面裂縫到廣告牌檢測最后落地的方案里 UNet 和它的變體占了大多數(shù)。這個資源提供了一個可以直接跑的 Python 實現(xiàn)適合剛接觸分割的入門者也適合需要快速做可行性驗證的從業(yè)者。它能幫你把一個圖像分割模型從數(shù)據(jù)準(zhǔn)備、訓(xùn)練到預(yù)測完整跑通整個鏈路不超過 200 行核心代碼。2. UNet 網(wǎng)絡(luò)結(jié)構(gòu)為什么所有分割任務(wù)都繞不開它2.1 編碼器-解碼器骨架和跳躍連接UNet 的結(jié)構(gòu)可以拆成三段看編碼器、解碼器和跳躍連接。編碼器就是不斷卷積加下采樣把輸入從 512×512 壓到 16×16這個過程中通道數(shù)從 3 漲到 512特征圖越來越抽象對應(yīng)的是“語義信息”。解碼器則反向操作一步步上采樣把特征圖恢復(fù)回原分辨率這時候網(wǎng)絡(luò)要學(xué)會把抽象語義和空間位置對應(yīng)起來。如果沒有跳躍連接解碼器在恢復(fù)細(xì)節(jié)時會嚴(yán)重丟信息因為下采樣把空間位置徹底打散了。跳躍連接做的事情非常樸素把編碼器某一層的特征圖直接拼到解碼器對應(yīng)層上。例如第 3 層下采樣前的輸出是 64×64×256解碼器在同一尺度上恢復(fù)時就把這兩個特征圖在通道維度上拼起來得到 64×64×512然后再卷積。這讓解碼器既能看到高層語義又能直接訪問淺層的邊緣紋理信息。實際訓(xùn)練中的感受是去掉跳躍連接后模型分割出來的物體邊緣會出現(xiàn)明顯的“膨脹”或“鋸齒”加上之后效果會立刻干凈很多。這個設(shè)計讓 UNet 在小數(shù)據(jù)集上表現(xiàn)特別好因為它不需要從零學(xué)習(xí)“邊緣在哪里”這件事情圖像的低層結(jié)構(gòu)已經(jīng)通過跳躍連接直接送到解碼器了。2.2 深度、通道數(shù)與輸入尺寸的選型UNet 的原始論文用的是 32 作為初始通道數(shù)每下采樣一次通道翻倍直到 512。這個配置并不是拍腦袋定的它的出發(fā)點(diǎn)是讓每一層都有足夠的表達(dá)能力去編碼對應(yīng)尺度的信息。實際使用時要看“你的物體有多大”這個關(guān)鍵問題。如果做的是車道線分割或裂縫檢測物體細(xì)長輸入尺寸建議至少 512×512初始通道數(shù)可以降到 16因為 32 通道在推理時會明顯拖慢速度而精度提升很有限。如果做的是地毯、廣告牌這類大塊面物體256×256 輸入加初始通道 16 已經(jīng)足夠硬上 512 輸入只會增加顯存占用。編碼器深度的選擇和一些常見的誤區(qū)有關(guān)。很多人以為層數(shù)越多效果越好但 UNet 下采樣到 32×32 以下時對大多數(shù)分割任務(wù)已經(jīng)過于抽象反而會讓小目標(biāo)直接消失。我在工程里常用的判斷標(biāo)準(zhǔn)是感受野至少要覆蓋目標(biāo)最大尺寸的一半。用 512×512 輸入、5 次下采樣的 UNet最后一層的感受野覆蓋整個圖但對 10×10 像素的小目標(biāo)來說它在下采樣第三次時就已經(jīng)被壓縮成了 1-2 個像素——所以遇到底層細(xì)節(jié)任務(wù)時我會減少下采樣次數(shù)或者跳過最后一個跳躍連接讓網(wǎng)絡(luò)在保留空間信息的前提下做特征融合比盲目加深有效得多。2.3 從 ResNet 編碼器到輕量化變體的改進(jìn)邏輯標(biāo)準(zhǔn) UNet 的編碼器是重復(fù)的卷積塊它的表達(dá)能力和 ResNet34 這類預(yù)訓(xùn)練骨干相比存在明顯差距。使用 ImageNet 預(yù)訓(xùn)練的 ResNet34 作為骨干遷移學(xué)習(xí)的效果幾乎總是優(yōu)于隨機(jī)初始化的標(biāo)準(zhǔn) UNet尤其當(dāng)訓(xùn)練數(shù)據(jù)只有幾百張圖時。原因很直觀預(yù)訓(xùn)練骨干已經(jīng)學(xué)到了通用的邊緣、紋理、顏色特征分割需要的底層視覺能力不需要從零學(xué)。代碼里把 UNet 的編碼器部分替換成 ResNet34解碼器保持原樣然后把跳躍連接的通道數(shù)改成 ResNet 各層的輸出通道即可。輕量化方向則是另一個極端對實時性要求高的場景MobilenetV2 作為編碼器的 UNet 是常見選擇。它用深度可分離卷積把計算量大幅降低。實際使用時要注意MobilenetV2 的特征圖通道排布和標(biāo)準(zhǔn) UNet 不同跳躍連接拼接時要讀一下各層的 out_channels而不是直接照搬源碼。至于注意力機(jī)制像在跳躍連接后面加一個簡單的 SE 模塊或 CBAM確實能提升邊界處的精度但代價是參數(shù)和推理耗時增加在工程里我通常只在最后一層跳躍連接上加而不是每層都加。3. 數(shù)據(jù)準(zhǔn)備與預(yù)處理決定效果的七成因素3.1 數(shù)據(jù)目錄組織與標(biāo)簽格式分割數(shù)據(jù)集的標(biāo)配是“圖像 掩膜”的組合。常見做法是按images和masks兩個目錄平鋪存放文件名一一對應(yīng)。掩膜有兩種格式一種是 PNG 索引圖像素值為 0、1、2 這樣的類別 ID另一種是灰度圖或 RGB 圖用 0 和 255 表示前景背景。這看起來是小事但數(shù)據(jù)加載代碼必須知道自己處理的是哪一種否則后面訓(xùn)練會在標(biāo)簽讀取時吃大虧。我的經(jīng)驗是先寫一個數(shù)據(jù)檢查函數(shù)把每張掩膜的unique值打出來看看。索引圖應(yīng)該是[0, 1, 2]這種連續(xù)小整數(shù)而掩膜圖可能是[0, 255]這時需要在加載時做一個mask // 255的操作把 255 映射回 1。還有一個常見的坑是某些標(biāo)注工具導(dǎo)出的掩膜帶有調(diào)色板模式用cv2.imread讀出來會是三通道。這時候要cv2.imread(path, cv2.IMREAD_GRAYSCALE)強(qiáng)制讀成灰度圖否則三通道的掩膜會被誤當(dāng)作 RGB 圖像參與訓(xùn)練損失函數(shù)計算時直接出錯。建議第一次跑通前先打印三張樣本的 shape 和 dtype 確認(rèn)。數(shù)據(jù)增強(qiáng)的順序也常有人搞反先對圖像做幾何變換再對掩膜做完全相同的變換最后再分別做歸一化。這個過程中最怕的是用torchvision.transforms里的隨機(jī)類變換它對圖像和掩膜是分別隨機(jī)化的導(dǎo)致圖像旋轉(zhuǎn)了 10 度而掩膜旋轉(zhuǎn)了 15 度。我的做法是自定義一個SegPairTransform類在同一個seed下生成隨機(jī)參數(shù)然后分別應(yīng)用到圖像和掩膜上保證它們空間變換完全一致。做分割項目時數(shù)據(jù)增強(qiáng)的 bug 往往是訓(xùn)練時 loss 正常下降、驗證時指標(biāo)卻一直上不去的頭號元兇因為它把標(biāo)簽錯位了卻不報錯。3.2 加載器與增強(qiáng)策略的完整代碼以下是我在分割項目里常用的數(shù)據(jù)加載器實現(xiàn)按“讀圖 → 統(tǒng)一尺寸 → 成對增強(qiáng) → 張量化”的流程走import cv2 import numpy as np import torch from torch.utils.data import Dataset from torchvision import transforms class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, size(512, 512), augmentFalse): self.img_paths sorted(list(img_dir.glob(*.png))) # 按文件名排序 self.mask_paths sorted(list(mask_dir.glob(*.png))) self.augment augment self.resize_h, self.resize_w size assert len(self.img_paths) len(self.mask_paths), 圖像和掩膜數(shù)量不一致 def __getitem__(self, idx): img cv2.imread(str(self.img_paths[idx])) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) mask cv2.imread(str(self.mask_paths[idx]), cv2.IMREAD_GRAYSCALE) mask mask // 255 # 255 映射到 1 img cv2.resize(img, (self.resize_w, self.resize_h), interpolationcv2.INTER_LINEAR) mask cv2.resize(mask, (self.resize_w, self.resize_h), interpolationcv2.INTER_NEAREST) if self.augment: img, mask self._pair_augment(img, mask) img torch.from_numpy(img).permute(2, 0, 1).float() / 255.0 mask torch.from_numpy(mask).long() return img, mask def _pair_augment(self, img, mask): # 同一個隨機(jī)種子保證空間變換一致 seed np.random.randint(0, 2**31) cv2.setRNGSeed(seed) if np.random.random() 0.5: img cv2.flip(img, 1) mask cv2.flip(mask, 1) cv2.setRNGSeed(seed) # 重置種子讓旋轉(zhuǎn)參數(shù)保持一致 angle np.random.uniform(-15, 15) h, w img.shape[:2] M cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1.0) img cv2.warpAffine(img, M, (w, h), flagscv2.INTER_LINEAR) mask cv2.warpAffine(mask, M, (w, h), flagscv2.INTER_NEAREST) return img, mask這段代碼里最關(guān)鍵的是mask // 255和重采樣方式。INTER_NEAREST是掩膜縮放時唯一推薦的方式因為它不會產(chǎn)生插值導(dǎo)致的偽類別比如背景經(jīng)過INTER_LINEAR后出現(xiàn) 128 這樣不屬于任何類別的中間值。圖像則用INTER_LINEAR保持邊緣平滑。旋轉(zhuǎn)使用的cv2.setRNGSeed是為了讓圖像和掩膜用同一組隨機(jī)參數(shù)否則分割標(biāo)簽會在空間上錯位。3.3 數(shù)據(jù)量不夠時怎么補(bǔ)救標(biāo)注數(shù)據(jù)不夠是分割項目的常態(tài)幾百張圖也同樣能訓(xùn)練出一個能用的模型。首要手段是離線增強(qiáng)除了在線增強(qiáng)的隨機(jī)翻轉(zhuǎn)旋轉(zhuǎn)再用 HSV 空間的隨機(jī)色度抖動增加顏色多樣性。另一個被驗證有效的做法是使用預(yù)訓(xùn)練編碼器把 UNet 的骨干換成 ResNet34 并用 ImageNet 權(quán)重初始化訓(xùn)練只需調(diào)解碼器部分幾百張圖完全足夠。還有一個方向是偽標(biāo)簽先用當(dāng)前模型對無標(biāo)簽數(shù)據(jù)做預(yù)測挑置信度高的區(qū)域作為訓(xùn)練數(shù)據(jù)——這個辦法在沒有標(biāo)注標(biāo)注人力時的確有效也直接幫助后續(xù)預(yù)測階段積累樣本。如果連幾百張圖都沒有就先別談訓(xùn)練直接把任務(wù)降級為目標(biāo)檢測或者用現(xiàn)成的分割模型先做一輪預(yù)標(biāo)注再人工修正效率遠(yuǎn)高于手工從頭標(biāo)注。4. 訓(xùn)練配置與計算損失函數(shù)、學(xué)習(xí)率和顯存管理4.1 損失函數(shù)的選擇邏輯UNet 最常用的損失函數(shù)組合是CrossEntropyLoss和DiceLoss兩類任務(wù)有明確的適用邊界。如果目標(biāo)是背景占比大而前景占比很小比如裂縫、血管CrossEntropyLoss會直接偏向背景類網(wǎng)絡(luò)預(yù)測結(jié)果可能全圖都是背景。DiceLoss 直接優(yōu)化前景區(qū)域的覆蓋率在這種類別極端不均衡的場景下效果更穩(wěn)定。如果目標(biāo)是前景和背景比例接近的比如地毯、廣告牌這種大塊物體BCEWithLogitsLoss就夠用配一個 0.2 到 0.4 的類別權(quán)重即可。工程里我一般不用單一的 DiceLoss而是用 BCE 和 Dice 加權(quán)相加經(jīng)驗值是BCE DiceLoss權(quán)重取 0.5 對 0.5。BCE 負(fù)責(zé)給所有像素提供穩(wěn)定的梯度信號DiceLoss 負(fù)責(zé)把前景區(qū)域的梯度變陡。訓(xùn)練初期 BCE 主導(dǎo)后期 DiceLoss 主導(dǎo)兩者互補(bǔ)收斂速度通常比任何單一損失都要快。另外注意這里不要過早引入 Focal Loss 這類復(fù)雜損失它會加重訓(xùn)練的不穩(wěn)定性先跑通流程再說。4.2 訓(xùn)練循環(huán)、學(xué)習(xí)率與全流程代碼核心訓(xùn)練循環(huán)我把關(guān)鍵部分寫出來包含梯度裁剪、學(xué)習(xí)率調(diào)度和驗證指標(biāo)。import torch import torch.nn as nn def train_one_epoch(model, loader, optimizer, criterion, device, clip1.0): model.train() total_loss 0.0 for images, masks in loader: images, masks images.to(device), masks.to(device) logits model(images) # 輸出 shape: [B, 1, H, W] loss criterion(logits, masks.float().unsqueeze(1)) optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(model.parameters(), clip) # 防梯度爆炸 optimizer.step() total_loss loss.item() return total_loss / len(loader) criterion lambda out, target: \ nn.BCEWithLogitsLoss()(out, target) 0.5 * dice_loss(out, target) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-5) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max60)dice_loss是自定義函數(shù)對輸出先過sigmoid然后計算1 - (2*交集 平滑項)/(并集 平滑項)。clip_grad_norm_是分割網(wǎng)絡(luò)容易忽略的一步UNet 解碼器層數(shù)多反向傳播時梯度有累積效應(yīng)尤其是輸入尺寸較大時不裁剪的話訓(xùn)練中期容易突然發(fā)散。學(xué)習(xí)率從1e-4開始是經(jīng)驗值太高會出現(xiàn) loss 震蕩太低則收斂過慢。CosineAnnealingLR配合 60 個 epoch 的調(diào)度被驗證效果穩(wěn)定。我一般在訓(xùn)練時同時保留兩份 checkpoint一份按驗證 IoU 最高的一份按最新的避免過擬合導(dǎo)致最佳模型丟失。4.3 顯存不足和批大小的心得顯存不足是新人第一個會撞到的墻。多數(shù)入門顯卡 6GB 顯存跑 512×512 輸入batch size 只能設(shè)到 2 甚至 1這不是很好但也能跑。優(yōu)先記住先確認(rèn)輸入尺寸是否真的需要 512如果是 256 就能覆蓋目標(biāo)那就用 256。真正對顯存影響最大的不是輸入的長寬而是跳躍連接里的特征圖4 個尺度的特征圖會在拼接時同時駐留顯存。減少初始通道數(shù)可以從 32 降到 16顯存立刻能省三分之一左右。訓(xùn)練時使用混合精度會再省一筆常見做法是在 PyTorch 里用自動混合精度。另一個被很多人忽略的點(diǎn)是驗證階段的顯存峰值。訓(xùn)練時梯度是逐層釋放的驗證時不需要梯度但整張?zhí)卣鲌D會保留到最后所以驗證時的顯存占用反而比訓(xùn)練時更穩(wěn)定。如果驗證時報 OOM說明是你的輸入尺寸或通道數(shù)實在太高要回頭降配置。5. 避坑排查UNet 使用時的常見問題5.1 掩膜尺寸和原圖對不上現(xiàn)象訓(xùn)練 loss 能下降但輸出圖明顯比輸入小一圈或者有一些奇怪的黑色邊緣。原因數(shù)據(jù)集的長寬不是固定比例被直接強(qiáng)行 resize 到正方形掩膜和原圖的內(nèi)容發(fā)生了非線性縮放語義錯位。另一類是數(shù)據(jù)增強(qiáng)時 train 和 val 用的尺寸不一致評估時輸入被打回原始尺寸導(dǎo)致輸出對不上。解決統(tǒng)一在數(shù)據(jù)加載階段用同一個(height, width)參數(shù)強(qiáng)轉(zhuǎn)train和val的尺寸保持一致的配置。圖像 resize 用cv2.INTER_LINEAR掩膜用cv2.INTER_NEAREST這兩者配錯導(dǎo)致的內(nèi)容錯位從訓(xùn)練日志上是看不出來的只有在可視化掩膜覆蓋在原圖上時才會暴露。5.2 類別不均衡導(dǎo)致背景“吃掉”前景現(xiàn)象驗證集上 IoU 很高比如 0.95但預(yù)測圖全黑或者小目標(biāo)全丟。原因背景像素占 99%前景占 1%模型學(xué)到了一個“預(yù)測全背景”的捷徑。此時CrossEntropyLoss的數(shù)值很小因為絕大多數(shù)像素都被正確分類了但 IoU 指標(biāo)因為前景為空而直接變成 0。解決換成BCEWithLogitsLoss加 DiceLoss 的組合給前景類別加權(quán)。另一個辦法是裁剪訓(xùn)練區(qū)域只看包含前景的局部切片再把切片的坐標(biāo)記錄下來預(yù)處理階段把比例拉均衡。我檢查這個問題的習(xí)慣是訓(xùn)練到一半時打印一遍每張預(yù)測圖的 unique 值如果超過九成都是同一個值基本就是類別失衡不是模型沒學(xué)好。5.3 學(xué)習(xí)率設(shè)置不當(dāng)訓(xùn)練發(fā)散現(xiàn)象loss 在訓(xùn)練前幾個 batch 正常下降到第 20 個 epoch 附近突然飆到 10 倍以上之后很難恢復(fù)。原因?qū)W習(xí)率太高或者余弦退火降得太快讓權(quán)重更新跨過了局部最優(yōu)而且 UNet 的跳躍連接會放大這種不穩(wěn)定。解決把lr從1e-3降到1e-4并加上nn.utils.clip_grad_norm_(model.parameters(), 1.0)。更穩(wěn)的做法是每 5 個 epoch 保存一次模型備份訓(xùn)練結(jié)束后選擇驗證指標(biāo)最佳的那份而不是選擇最后一份。不要過度依賴早停分割任務(wù)的驗證指標(biāo)波動比分類任務(wù)大很多早停閾值設(shè)置過小會把還有上升空間的模型直接砍掉。5.4 數(shù)據(jù)增強(qiáng)把掩膜“旋轉(zhuǎn)錯位”之后模型肉眼看不太出來但指標(biāo)不變現(xiàn)象訓(xùn)練 loss 下降正常驗證 IoU 停滯不前多個 epoch 都一樣。原因圖像做了旋轉(zhuǎn)而掩膜沒有或者用了不同的隨機(jī)種子導(dǎo)致模型始終學(xué)到“把分割結(jié)果偏移某個角度”的錯位映射。肉眼檢查單張圖有時候看不出問題這里的關(guān)鍵是統(tǒng)計驗證集所有樣本的 IoU如果多類別任務(wù)某些類別 IoU 一直為 0 但整體指標(biāo)尚可說明增強(qiáng)和驗證邏輯之間存在不一致。解決統(tǒng)一用成對變換作為修復(fù)手段我在__getitem__里把掩膜直接在數(shù)據(jù)加載階段打印出來和原圖疊加可視化人工過一遍確認(rèn)每個通道都逐像素對齊后再開始訓(xùn)練。這類問題和大多數(shù)框架報錯無關(guān)屬于“沉默失敗”依賴可視化檢查屬于血淚教訓(xùn)。6. 推理與可視化把模型輸出變成能看的結(jié)果模型訓(xùn)練完只是第一步真正要交付的是“輸入任意一張圖 → 輸出帶標(biāo)注的圖”。推理階段的核心是把模型的 logits 轉(zhuǎn)成類別掩膜再疊加回原圖。這里有一個容易出問題的細(xì)節(jié)預(yù)測時要保持和訓(xùn)練時相同的預(yù)處理流程尤其是mean/std歸一化和 resize 尺寸否則預(yù)測效果會有肉眼可見的下降。我常用的推理代碼是把模型的sigmoid輸出以 0.5 為閾值生成二值掩膜然后用半透明的紅色在原圖上疊加再用opencv的addWeighted混合。保存結(jié)果時用matplotlib或cv2.imwrite直接存成 PNG。如果同時需要計算 IoU要記得把掩膜和標(biāo)簽都轉(zhuǎn)成torch.int64并壓縮到[0, 1]避免數(shù)組廣播帶來的尺寸錯位問題。整套流程跑通之后還有兩件事值得做。一件是把模型導(dǎo)出成 TorchScript 或 ONNX這樣可以在移動端和邊緣設(shè)備上做推理省去模型架構(gòu)和預(yù)處理代碼的重復(fù)編寫開銷我的經(jīng)驗是 ONNX 導(dǎo)出最省心直接用torch.onnx.export配合固定尺寸輸入就能一步到位。另一件事是把數(shù)據(jù)增強(qiáng)、訓(xùn)練和推理腳本封裝成一個帶命令行參數(shù)的main.py按數(shù)據(jù)路徑和模型權(quán)重兩個參數(shù)就能調(diào)起整個流程這不算復(fù)雜工程但對復(fù)現(xiàn)和后續(xù)修改的幫助非常大。從那以后我每次做分割項目都會在訓(xùn)練前強(qiáng)制走一遍“數(shù)據(jù)可視化 → 跑通 1 個 batch → 再加增強(qiáng)”的流程這三個步驟能攔截掉大部分靜默 bug。希望這篇實戰(zhàn)筆記對你跑通 UNet 圖像分割有幫助。本文還有配套的精品資源點(diǎn)擊獲取