【语义分割系列】ICNET(实时语义分割)理论以及代码实现

git地址:https://github.com/Tramac/awesome-semantic-segmentation-pytorch

包括:

FCN
ENet
PSPNet
ICNet
DeepLabv3
DeepLabv3+
DenseASPP
EncNet
BiSeNet
PSANet
DANet
OCNet
CGNet
ESPNetv2
CCNet
DUNet(DUpsampling)
FastFCN(JPU)
LEDNet
Fast-SCNN
LightSeg
DFANet

几乎囊括了语义分割的主流源码.调试了一个效果还不错

看自己需要选择吧。

ICNET理论的几个关键点:

1.核心思想:

        将图像分为高中低三层。性能大幅提升的根本原因就是让低低分辨的图像经过语义分割网络产生粗糙的分割结果;之后特征级联混合单元(cascade label guidance)与标签引导的级联策略(cascade label guidance strategy)将中分辨率和高分辨率的特征整合,逐步地优化之前生成的粗糙分割结果。

2.网络图:

        由下至上为高到低分辨率,低分辨率的结果通过CFF与中分辨率结果组合生成蓝色方块,中分辨率结果与蓝色方块通过CFF方法与高分辨率特征图融合。最终将结果上采样获得最终输出。

3.CFF:

      Cascade Feature Fusion,

首先对F1做双线性插值,2倍上采样,使其与F2大小同

然后用C3X3,采样率为2的空洞卷积对F1做上采样。生成图像大小与F2同。空洞卷积结合了原始相邻像素特征信息,感受野与F2同。计算量小于反卷积。记为A

对特征F2 使用CX1X1的卷积将其与F1通道相同(相当于升维)。记为B

使用BN层标准化两个处理好的特征层 A B

把二者加起来在用relu激活 得到与F2大小一致的融合特征F2'

4.Cascade Label Guidance

        红圈4的部分,这个是在训练时候才用,对应下面的代码的x_24_cls部分。这样的 位置一共有三个。说白了就是提前级联的地方,不同分辨率下的标注图的标注信息

x_cff_24, x_24_cls = self.cff_24(x_sub4, x_sub2)

 

5.性能:

我主要想用实时语义分割。就随便选了个ICNET试验。2048 *1024的图片预测GPU预测时间190ms左右。如果改改图片尺寸,网络参数再压缩一下 能快不少

6.代码:

       注意1:我测试选的主干网络选的是resnet50,需要对着这个网络看下面的注释。 才能一目了然   主干网络的作用就是提取特征而已。其他主干网络类似,主要为了看通顺ICNET实现流程而已。

       注意2: 代码CascadeFeatureFusion部分有一处不太明白的地方。看论文这句 

应该是与fig.3对应的。与我理解一致,应该是直接对low图做完上采样后,直接拿上采样作为输入,类别数作为输出通道做了次卷积。所以对其中做了下修改,还没训练完事,等训练完事看效果吧。

class CascadeFeatureFusion(nn.Module):
    """CFF Unit"""


        self.conv_low_cls = nn.Conv2d(low_channels, nclass, 1, bias=False)


        x_low = F.interpolate(x_low, size=x_high.size()[2:], mode='bilinear', align_corners=True)
        x_low1 = x_low
        
        x_low_cls = self.conv_low_cls(x_low1)

        return x, x_low_cls

   完整代码:                                                

"""Image Cascade Network"""
import torch
import torch.nn as nn
import torch.nn.functional as F

from .segbase import SegBaseModel

__all__ = ['ICNet', 'get_icnet', 'get_icnet_resnet50_citys',
           'get_icnet_resnet101_citys', 'get_icnet_resnet152_citys']

class ICNet(SegBaseModel):
    """Image Cascade Network"""

    def __init__(self, nclass, backbone='resnet50', aux=False, jpu=False, pretrained_base=False, **kwargs):
        super(ICNet, self).__init__(nclass, aux, backbone, pretrained_base=pretrained_base, **kwargs)
        self.conv_sub1 = nn.Sequential(
            _ConvBNReLU(3, 32, 3, 2, **kwargs),#ksize 3  stride 2   
            _ConvBNReLU(32, 32, 3, 2, **kwargs),
            _ConvBNReLU(32, 64, 3, 2, **kwargs)
        )

        self.ppm = PyramidPoolingModule()

        self.head = _ICHead(nclass, **kwargs)

        self.__setattr__('exclusive', ['conv_sub1', 'head'])

    def forward(self, x):
        # sub 1  原图做普通卷积 BN RL                     对应fig.2的 红圈1    例如 x shape [20, 3, 480, 480]
        x_sub1 = self.conv_sub1(x)  #图像变为原图 1/8                               x_sub1 shape [20, 64, 60, 60]

        # sub 2  x_sub2是将原图缩放到  1/2                                          x_sub2 shape [20, 3, 240, 240]
        x_sub2 = F.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=True)
        #过主干网络 例如resnet之类。 对x_sub2做特征提取   对应fig.2的 红圈2  看取值是取的layer2   大小变为  1 / 16
        #                                                                           x_sub2 shape [20, 512, 30, 30]
        _, x_sub2, _, _ = self.base_forward(x_sub2)

        # sub 4 x_sub4 原图缩放到 1 / 4                                             x_sub4 shape [20, 3, 120, 120]
        x_sub4 = F.interpolate(x, scale_factor=0.25, mode='bilinear', align_corners=True)#双线性插值将为 原图 1/4
        #过主干网络  对x_sub4做特征提取                   对应fig.2 的红圈3   看取值是取的layer4          1/32
        #                                                                           x_sub4 shape [20, 2048, 15, 15]
        _, _, _, x_sub4 = self.base_forward(x_sub4)
        # add PyramidPoolingModule
        # x_sub4作为输出 就是加上全局池化后的信息了                                 x_sub4 shape [20, 2048, 15, 15]
        x_sub4 = self.ppm(x_sub4)
        #上采样部分  x_sub1 [20, 64, 60, 60]  x_sub2 [20, 512, 30, 30]  x_sub4 [20, 2048, 15, 15]
        outputs = self.head(x_sub1, x_sub2, x_sub4)

        return tuple(outputs)
#类似PPM结构
class PyramidPoolingModule(nn.Module):
    def __init__(self, pyramids=[1,2,3,6]):
        super(PyramidPoolingModule, self).__init__()
        self.pyramids = pyramids

    def forward(self, input):
        feat = input
        height, width = input.shape[2:]
        for bin_size in self.pyramids:
            #池化 x的尺寸分别为 1*1  2*2 3*3 6*6 通道数不变 与输入特征图一致
            #这么做点目的是保留全局信息,防止大块的目标中出现空洞,大目标检测不好 例如 马路 例如 天空
            x = F.adaptive_avg_pool2d(input, output_size=bin_size)
            # 将池化结果插值回原图大小
            x = F.interpolate(x, size=(height, width), mode='bilinear', align_corners=True)
            # 输入特征图与插值图像叠加
            feat  = feat + x
        return feat

class _ICHead(nn.Module):
    def __init__(self, nclass, norm_layer=nn.BatchNorm2d, **kwargs):
        super(_ICHead, self).__init__()
        self.cff_12 = CascadeFeatureFusion(512, 64, 128, nclass, norm_layer, **kwargs)
        #self.cff_12 = CascadeFeatureFusion(128, 64, 128, nclass, norm_layer, **kwargs)  
        self.cff_24 = CascadeFeatureFusion(2048, 512, 128, nclass, norm_layer, **kwargs)

        self.conv_cls = nn.Conv2d(128, nclass, 1, bias=False)

    def forward(self, x_sub1, x_sub2, x_sub4):
        outputs = list()
        # cff_24是论文中的CFF部分实现代码 这个是做最小和次小的特征融合呢
        # x_cff_24 shape [20, 128, 30, 30]   x_24_cls shape [20, 19, 30, 30]
        x_cff_24, x_24_cls = self.cff_24(x_sub4, x_sub2)
        outputs.append(x_24_cls)
        #x_cff_12, x_12_cls = self.cff_12(x_sub2, x_sub1)
        # 论文中的CFF部分实现代码 这个是做次小和最大的特征融合呢
        # # x_cff_12 shape [20, 128, 60, 60]   x_12_cls shape [20, 19, 60, 60]
        x_cff_12, x_12_cls = self.cff_12(x_cff_24, x_sub1)  
        outputs.append(x_12_cls)
        # x_cff_12 上采样一次           up_x2 shape [20, 128, 120, 120]
        up_x2 = F.interpolate(x_cff_12, scale_factor=2, mode='bilinear', align_corners=True)
        # 做一次loss函数需要的卷积       up_x2 shape [20, 19, 120, 120]     
        up_x2 = self.conv_cls(up_x2)
        outputs.append(up_x2)
        # up_x2 再做一次上采样  直接放大到原图大小
        # x_cff_12 上采样一次  up_x2 shape [20, 19, 480, 480]
        up_x8 = F.interpolate(up_x2, scale_factor=4, mode='bilinear', align_corners=True)
        outputs.append(up_x8)
        # 1 -> 1/4 -> 1/8 -> 1/16  将list反序
        outputs.reverse()
        # outputs 依次存放{[20, 19, 480, 480], [20, 19, 120, 120], [20, 19, 60, 60], [20, 19, 30, 30]}
        return outputs


class _ConvBNReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1,
                 groups=1, norm_layer=nn.BatchNorm2d, bias=False, **kwargs):
        super(_ConvBNReLU, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
        self.bn = norm_layer(out_channels)
        self.relu = nn.ReLU(True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x


class CascadeFeatureFusion(nn.Module):
    """CFF Unit"""

    def __init__(self, low_channels, high_channels, out_channels, nclass, norm_layer=nn.BatchNorm2d, **kwargs):
        super(CascadeFeatureFusion, self).__init__()
        self.conv_low = nn.Sequential(
            nn.Conv2d(low_channels, out_channels, 3, padding=2, dilation=2, bias=False),
            norm_layer(out_channels)
        )
        self.conv_high = nn.Sequential(
            nn.Conv2d(high_channels, out_channels, 1, bias=False),
            norm_layer(out_channels)
        )
        self.conv_low_cls = nn.Conv2d(out_channels, nclass, 1, bias=False)
    # 对应论文fig.3 
    def forward(self, x_low, x_high):
        # 先将 x_low特征图做双线性插值上采样                    对应fig.3 图中框1   
        # 例如对于cff_24   x_low 输入为[20, 2048, 15, 15] 输出为[20, 2048, 30, 30]
        x_low = F.interpolate(x_low, size=x_high.size()[2:], mode='bilinear', align_corners=True)
        # 再对x_low特征图做空洞卷积     结果通道为128           对应fig.3 图中框2
        # 例如对于cff_24   x_low 输入为[20, 2048, 30, 30] 输出为[20, 128, 30, 30]
        x_low = self.conv_low(x_low)
        # 对x_high特征图做 1*1的卷积    结果通道为128           对应fig.3 图中框3
        # 例如对于cff_24   x_high 输入为[20, 512, 30, 30] 输出为[20, 128, 30, 30]
        x_high = self.conv_high(x_high)
        # 求和做relu 作为输出数据之一                           对应fig.3 图中框4
        x = x_low + x_high
        x = F.relu(x, inplace=True)
        # 对空洞x_low特征图做直接做卷积并做BatchNorm2d 结果通道数为类别维度   用来算损失函数的  对应fig.3 图中 框5  
        # 代码与论文图对应不上了   按理应该是双线性插值的结果作为输入才对  
        # 例如对于cff_24   x_low 输入为[20, 128, 30, 30] 输出为[20, 19, 30, 30]
        # 按论文画的图 应该为  x_low 输入为[20, 2048, 30, 30] 输出为[20, 19, 30, 30]才对  改了训练一次看看效果  这块是
        # 这个代码我没看懂的地方之一
        x_low_cls = self.conv_low_cls(x_low)
        return x, x_low_cls


def get_icnet(dataset='citys', backbone='resnet50', pretrained=False, root='~/.torch/models',
              pretrained_base=False, **kwargs):
    acronyms = {
        'pascal_voc': 'pascal_voc',
        'pascal_aug': 'pascal_aug',
        'ade20k': 'ade',
        'coco': 'coco',
        'citys': 'citys',
    }
    from ..data.dataloader import datasets
    model = ICNet(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
    path = '/home/yuany/awesome-semantic-segmentation-pytorch-master/scripts/models/best.pth'
    pretrained = True
    if pretrained:
        from .model_store import get_model_file
        #device = torch.device(kwargs['local_rank'])
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        #model.load_state_dict(torch.load(get_model_file('icnet_%s_%s' % (backbone, acronyms[dataset]), root=root),
        #                      map_location=device))
        model.load_state_dict(torch.load(path,
                              map_location=device))
    return model


def get_icnet_resnet50_citys(**kwargs):
    return get_icnet('citys', 'resnet50', **kwargs)


def get_icnet_resnet101_citys(**kwargs):
    return get_icnet('citys', 'resnet101', **kwargs)


def get_icnet_resnet152_citys(**kwargs):
    return get_icnet('citys', 'resnet152', **kwargs)


if __name__ == '__main__':
    img = torch.randn(1, 3, 256, 256)
    model = get_icnet_resnet50_citys()
    outputs = model(img)

再记录个新的更快的。先把图存着 后续看代码 看论文

完整步骤:我主要想做道路上的语义分割,选择的数据集是cityscapes。

1.准备数据集 cityscapes。把需要解压的压缩包解压一下就好

2.准备代码,github上的链接https://github.com/Tramac/awesome-semantic-segmentation-pytorch

3.解压后进入scripts文件夹,里面,有train文件。core文件夹存的是model。各种语义分割的模型都在里面了,在train的参数列表可见选择项目。

4.在cityscapes.py中修改下路径:直接把数据集cityscapes放在datasets路径下面了 ../datasets/cityscapes

5.datasets里面的citys文件我也改了下,放的是路径  例如我的路径是;/home/yuany/awesome-semantic-segmentation-pytorch-master/datasets/cityscapes 没注意这两个路径哪个起作用

6.然后就是输入train命令了 例如我model 选择 icnet  主干网络选resnet50  数据集选citys   初始学习率选0.01  迭代次数选择30000 那么命令如下:

python train.py --model icnet --backbone resnet50 --dataset citys --lr 0.01 --epochs 30000

7.训练完的model在models文件夹内,有一个最好的 还有一个最新的

8.使用:直接按demo去使用即可,设置好待分割图像的路径。分割即可,我是在icnt中设置的载入模型路径,就是图方便了

猜你喜欢

转载自blog.csdn.net/gbz3300255/article/details/111545273
今日推荐