前言
因為畢設中的部分內(nèi)容涉及到衛(wèi)星遙感影像道路分割，因此去對相關(guān)算法做了一些調(diào)研。
本文所使用數(shù)據(jù)集為DeepGlobe，來自于CVPR2018年的一個挑戰(zhàn)賽：DeepGlobe Road Extraction Challenge。
D-LinkNet為該挑戰(zhàn)賽的冠軍算法。
考慮到D-LinkNet開發(fā)版本較老(Python 2.7、Pytorch 0.2.0)，我對此項目進行了重構(gòu)，具體工作如下：
- 修改相關(guān)Python2語法，以滿足Python3.8開發(fā)環(huán)境
- 移除多卡訓練部分(
DataParallel)，以便讓代碼變得更加清晰易讀 - 增加模型驗證函數(shù)(
eval.py)，增加mIou指標以評估模型效果 - 增加新算法NL-LinkNet，并提供相關(guān)訓練結(jié)果
目前該倉庫支持下列分割算法：
- UNet
- D-UNet
- LinkNet
- D-LinkNet
- NL-LinkNet
項目地址：https://github.com/zstar1003/Road-Extraction
DeepGlobe數(shù)據(jù)集簡介
DeepGlobe數(shù)據(jù)集下載地址：https://pan.baidu.com/s/1chOnMUIzcKUzQr1LpuJohw?pwd=8888
該數(shù)據(jù)集包含6226張訓練圖片，每張圖片尺寸為1024×1024，圖像分辨率為0.5米/像素
數(shù)據(jù)預覽：
D-LinkNet網(wǎng)絡結(jié)構(gòu)
圖像分割在衛(wèi)星遙感道路分割領(lǐng)域大致有以下一系列算法，算法發(fā)布時間線如下：
FCN(2015)->UNet(2015) -> LinkNet(2017)->D-LinkNet(2018)->NL-LinkNet(2019)->…
D-LinkNet的網(wǎng)絡結(jié)構(gòu)如下圖所示：
這個網(wǎng)絡整體結(jié)構(gòu)和UNet比較類似，主要在此架構(gòu)中加了一些小改進，如殘差塊、空洞卷積等。改進提升比較明顯的是該算法引入了TTA(Test Time Augmentation)策略，即測試時加強，后面將對此進行詳解。
修改模型結(jié)構(gòu)層名
由于我移除了DataParallel多卡并行訓練的結(jié)構(gòu)，直接加載官方提供的模型會報錯：
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”
…
這是由于模型結(jié)構(gòu)層名不一致，模型文件中包含的層名多了module.，因此寫了個腳本進行轉(zhuǎn)換utils/turn_model.py：
import collections
import torch
if __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的思想就是在測試時使用數(shù)據(jù)增強，比如一張圖片直接進行分割，得到的效果可能有限，那么將這副圖片進行旋轉(zhuǎn)、翻轉(zhuǎn)等數(shù)據(jù)增強方式，進行分割，最后將所有分割結(jié)果進行疊加。
下面來按程序運行邏輯的順序進行分析：
首先，程序加載完一張圖片后，img是原圖，img90是將圖像逆時針旋轉(zhuǎn)90度，相關(guān)代碼：
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])
之后，構(gòu)建了一個img2，在img1的第二個維度進行逆序，實現(xiàn)垂直翻轉(zhuǎn)
img2 = np.array(img1)[:, ::-1] # 垂直翻轉(zhuǎn)
直觀顯示：
- show_img(img2[0], img2[1])
img3同理，在img1的第三個維度進行逆序，實現(xiàn)水平翻轉(zhuǎn)
img3 = np.array(img1)[:, :, ::-1] # 水平翻轉(zhuǎn)
直觀顯示：
- show_img(img3[0], img3[1])
img4是對img2的實現(xiàn)水平翻轉(zhuǎn)，等價于對img1進行水平和垂直翻轉(zhuǎn)
img4 = np.array(img2)[:, :, ::-1] # 垂直翻轉(zhuǎn)+水平翻轉(zhuǎn)
直觀顯示：
- show_img(img4[0], img4[1])
后面就是對每一個部分進行推理，然后最后返回的mask2是疊加后的結(jié)果，maska[0]是原始圖像的推理結(jié)果
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]
直觀進行比較，左側(cè)是原圖推理，右側(cè)是TTA后的推理結(jié)果：
- show_img(maska[0], mask2)
可以看到，使用TTA的效果還是挺明顯的。
NL-LinkNet
2019年，NL-LinkNet被提出，據(jù)稱，它在DeepGlobe數(shù)據(jù)集上mIOU高于D-LinkNet.
相關(guān)倉庫：https://github.com/yswang1717/NLLinkNet
由于倉庫作者提供的模型推理效果很差(可能作者傳錯了文件)，我又在自己的RTX2060上訓練了128epoch（實際設置200個epoch，128個epoch之后模型收斂早停)。模型訓練起來還是比較慢的，耗費時間約57小時，具體日志信息可參看logs。文章來源：http://www.zghlxwxcb.cn/news/detail-567903.html
下面是兩個模型對同一幅圖片的分割結(jié)果比較：
可以看到，NL-LinkNet分割結(jié)果更加順滑一些。文章來源地址http://www.zghlxwxcb.cn/news/detail-567903.html
到了這里，關(guān)于【圖像分割】衛(wèi)星遙感影像道路分割：D-LinkNet算法解讀的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
