【图像分割】卫星遥感影像道路分割：D-LinkNet算法解读
前言
因為畢設中的部分內容涉及到衛星遙感影像道路分割，因此去對相關算法做了一些調研。
本文所使用數據集為DeepGlobe，來自于CVPR2018年的一個挑戰賽：DeepGlobe Road Extraction Challenge。
D-LinkNet為該挑戰賽的冠軍算法。
考慮到D-LinkNet開發版本較老(Python 2.7、Pytorch 0.2.0)，我對此項目進行了重構，具體工作如下：
- 修改相關Python2語法，以滿足Python3.8開發環境
- 移除多卡訓練部分(DataParallel)，以便讓代碼變得更加清晰易讀
- 增加模型驗證函數(eval.py)，增加mIou指標以評估模型效果
- 增加新算法NL-LinkNet，并提供相關訓練結果
目前該倉庫支持下列分割算法：
- UNet
- D-UNet
- LinkNet
- D-LinkNet
- NL-LinkNet
項目地址：https://github.com/zstar1003/Road-Extraction
DeepGlobe數據集簡介
DeepGlobe數據集下載地址：https://pan.baidu.com/s/1chOnMUIzcKUzQr1LpuJohw?pwd=8888
該數據集包含6226張訓練圖片，每張圖片尺寸為1024×1024，圖像分辨率為0.5米/像素
數據預覽：
D-LinkNet網絡結構
圖像分割在衛星遙感道路分割領域大致有以下一系列算法，算法發布時間線如下：
FCN(2015)->UNet(2015) -> LinkNet(2017)->D-LinkNet(2018)->NL-LinkNet(2019)->…
D-LinkNet的網絡結構如下圖所示：
這個網絡整體結構和UNet比較類似，主要在此架構中加了一些小改進，如殘差塊、空洞卷積等。改進提升比較明顯的是該算法引入了TTA(Test Time Augmentation)策略，即測試時加強，后面將對此進行詳解。
修改模型結構層名
由于我移除了DataParallel多卡并行訓練的結構，直接加載官方提供的模型會報錯：
RuntimeError: Error(s) in loading state_dict for DinkNet34:
Missing key(s) in state_dict: “firstconv.weight”, “firstbn.weight”, “firstbn.bias”,
Unexpected key(s) in state_dict: “module.firstconv.weight”, “module.firstbn.weight”, “module.firstbn.bias”
…
這是由于模型結構層名不一致，模型文件中包含的層名多了module.，因此寫了個腳本進行轉換utils/turn_model.py：
import collections import torchif __name__ == '__main__':path = '../weights/log01_dink34.th'model = torch.load(path)new_model = collections.OrderedDict([(k[7:], v) if k[:7] == 'module.' else (k, v) for k, v in model.items()])torch.save(new_model, "../weights/dlinknet.pt")TTA策略
TTA的思想就是在測試時使用數據增強，比如一張圖片直接進行分割，得到的效果可能有限，那么將這副圖片進行旋轉、翻轉等數據增強方式，進行分割，最后將所有分割結果進行疊加。
下面來按程序運行邏輯的順序進行分析：
首先，程序加載完一張圖片后，img是原圖，img90是將圖像逆時針旋轉90度，相關代碼：
def segment(self, path):img = cv2.imread(path)img = cv2.resize(img, resize_settings) # Shape:(1024, 1024, 3)img90 = np.array(np.rot90(img)) # Shape:(1024, 1024, 3)img1 = np.concatenate([img[None, ...], img90[None, ...]]) # Shape:(2, 1024, 1024, 3) img[None]:增加第一個位置維度img1是將這兩張圖片拼接起來，下面直觀進行顯示查看：
- show_img(img1[0], img1[1])
之后，構建了一個img2，在img1的第二個維度進行逆序，實現垂直翻轉
img2 = np.array(img1)[:, ::-1] # 垂直翻轉直觀顯示：
- show_img(img2[0], img2[1])
img3同理，在img1的第三個維度進行逆序，實現水平翻轉
直觀顯示：
- show_img(img3[0], img3[1])
img4是對img2的實現水平翻轉，等價于對img1進行水平和垂直翻轉
img4 = np.array(img2)[:, :, ::-1] # 垂直翻轉+水平翻轉直觀顯示：
- show_img(img4[0], img4[1])
后面就是對每一個部分進行推理，然后最后返回的mask2是疊加后的結果，maska[0]是原始圖像的推理結果
maska = self.net.forward(img1).squeeze().cpu().data.numpy() # img1:Shape:(2, 1, 1024, 1024) -> (2, 1024, 1024) maskb = self.net.forward(img2).squeeze().cpu().data.numpy() maskc = self.net.forward(img3).squeeze().cpu().data.numpy() maskd = self.net.forward(img4).squeeze().cpu().data.numpy()mask1 = maska + maskb[:, ::-1] + maskc[:, :, ::-1] + maskd[:, ::-1, ::-1] mask2 = mask1[0] + np.rot90(mask1[1])[::-1, ::-1]直觀進行比較，左側是原圖推理，右側是TTA后的推理結果：
- show_img(maska[0], mask2)
可以看到，使用TTA的效果還是挺明顯的。
NL-LinkNet
2019年，NL-LinkNet被提出，據稱，它在DeepGlobe數據集上mIOU高于D-LinkNet.
相關倉庫：https://github.com/yswang1717/NLLinkNet
由于倉庫作者提供的模型推理效果很差(可能作者傳錯了文件)，我又在自己的RTX2060上訓練了128epoch（實際設置200個epoch，128個epoch之后模型收斂早停)。模型訓練起來還是比較慢的，耗費時間約57小時，具體日志信息可參看logs。
下面是兩個模型對同一幅圖片的分割結果比較：
可以看到，NL-LinkNet分割結果更加順滑一些。
總結
以上是生活随笔為你收集整理的【图像分割】卫星遥感影像道路分割：D-LinkNet算法解读的全部內容，希望文章能夠幫你解決所遇到的問題。
