前言
因为毕设中的部分内容涉及到卫星遥感影像道路分割,因此去对相关算法做了一些调研。
本文所使用数据集为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
。文章来源:https://www.toymoban.com/news/detail-567903.html
下面是两个模型对同一幅图片的分割结果比较:
可以看到,NL-LinkNet分割结果更加顺滑一些。文章来源地址https://www.toymoban.com/news/detail-567903.html
到了这里,关于【图像分割】卫星遥感影像道路分割:D-LinkNet算法解读的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!