前言
因为毕设中的部分内容涉及到卫星遥感影像道路分割，因此去对相关算法做了一些调研。
本文所使用数据集为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 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的思想就是在测试时使用数据增强，比如一张图片直接进行分割，得到的效果可能有限，那么将这副图片进行旋转、翻转等数据增强方式，进行分割，最后将所有分割结果进行叠加。
下面来按程序运行逻辑的顺序进行分析：
首先，程序加载完一张图片后，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的第三个维度进行逆序，实现水平翻转
img3 = np.array(img1)[:, :, ::-1] # 水平翻转
直观显示：
- 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分割结果更加顺滑一些。