【图像分割】卫星遥感影像道路分割：D-LinkNet算法解读

2023-07-01

模型 self img

前言因为毕设中的部分内容涉及到卫星遥感影像道路分割，因此去对相关算法做了一些调研。本文所使用数据集为DeepGlobe，来自于CVPR2018年的一个挑战赛：DeepGlobeRoadExtractionChallenge。D-LinkNet为该挑战赛的冠军算法。考虑到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 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")
1
2
3
4
5
6
7
8

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]:增加第一个位置维度
1
2
3
4
5

img1是将这两张图片拼接起来，下面直观进行显示查看：

show_img(img1[0], img1[1])

之后，构建了一个img2，在img1的第二个维度进行逆序，实现垂直翻转

img2 = np.array(img1)[:, ::-1]  # 垂直翻转
1

直观显示：

show_img(img2[0], img2[1])

img3同理，在img1的第三个维度进行逆序，实现水平翻转

img3 = np.array(img1)[:, :, ::-1]  # 水平翻转
1

直观显示：

show_img(img3[0], img3[1])

img4是对img2的实现水平翻转，等价于对img1进行水平和垂直翻转

img4 = np.array(img2)[:, :, ::-1]  # 垂直翻转+水平翻转
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]
1
2
3
4
5
6
7