语义分割——SegNet

南太湖小蚂蚁2024-11-17 8:16

SegNet是由剑桥大学团队开发的一个图像分割的开源项目，该项目可以对图像中的物体所在区域进行分割，例如车、马路、行人等，并且精确到像素级别。SegNet提出了一种编码器，解码器的结构，其实有点类似于FCN，但又有所不同。他的主要流程如下：

输入一幅待分割的图像，先进入编码器，再进入解码器，最后通过一个softmax得到每个像素的分类结果，也就是语义分割的结果。在我看来，SegNet和FCN最大的不同，也就是SegNet最大的特点就是它存储了编码过程中最大池化的索引。在SegNet网络结构中，进行2×2最大池化时，会存储相应的最大池化索引（位置）。在解码器处，执行上采样和卷积时，会调用相应编码器层处的最大池化索引以进行上采样。这种方式可以一定程度上解决物体边界划分不清的问题，因为上采样的信息是直接从原始输入图像中获取的，能够更准确地反映物体的边界。而FCN在上采样过程中，并没有考虑到编码时最大池化的索引位置，如下图所示：

整个SegNet的结构如下：

可以看到，编码器和解码器都有五个模块构成。

编码器1：两个卷积模块和一个最大池化模块（每个卷积模块包含一次卷积一次批归一化和一次非线性变换），大小缩小一半

编码器2：两个卷积模块和一个最大池化模块，大小缩小一半

编码器3：三个卷积模块和一个最大池化模块，大小缩小一半

编码器4：三个卷积模块和一个最大池化模块，大小缩小一半

编码器5：三个卷积模块和一个最大池化模块，大小缩小一半

解码器1：一个上采样模块和三个卷积模块，大小扩大一倍（在上采样过程中，使用保存的编码器最大池化时的索引）

解码器2：一个上采样模块和三个卷积模块，大小扩大一倍

解码器3：一个上采样模块和三个卷积模块，大小扩大一倍

解码器4：一个上采样模块和两个卷积模块，大小扩大一倍

解码器5：一个上采样模块和两个卷积模块，再拼接上一个softmax操作进行分类，大小扩大一倍，恢复成原始图像大小。

下面我们来看一下根据这个设计编写的代码：

python 复制代码 
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models

class SegNet(nn.Module):
    def __init__(self, num_classes=21):
        super(SegNet, self).__init__()

        self.encoder1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )
        self.encoder2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
        )
        self.encoder3 = nn.Sequential(
            nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )
        self.encoder4 = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
        )
        self.encoder5 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
        )

        # Decoder
        self.decoder1 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True)
        )
        self.decoder2 = nn.Sequential(
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )
        self.decoder3 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
        )
        self.decoder4 = nn.Sequential(
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.decoder5 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, num_classes, kernel_size=1)
        )

    def forward(self, x):
        # 用来保存各层的池化索引
        pool_indices = []
        x = self.encoder1(x)
        x, pool_indices1 = nn.MaxPool2d(2, stride=2, return_indices=True)(x)
        pool_indices.append(pool_indices1)
        print("x.shape: ",x.shape)
        x = self.encoder2(x)
        x, pool_indices2 = nn.MaxPool2d(2, stride=2, return_indices=True)(x)
        pool_indices.append(pool_indices2)
        print("x.shape: ",x.shape)
        x = self.encoder3(x)
        x, pool_indices3 = nn.MaxPool2d(2, stride=2, return_indices=True)(x)
        pool_indices.append(pool_indices3)
        print("x.shape: ",x.shape)
        x = self.encoder4(x)
        x, pool_indices4 = nn.MaxPool2d(2, stride=2, return_indices=True)(x)
        pool_indices.append(pool_indices4)
        print("x.shape: ",x.shape)
        x = self.encoder5(x)
        x, pool_indices5 = nn.MaxPool2d(2, stride=2, return_indices=True)(x)
        pool_indices.append(pool_indices5)
        print("x.shape: ",x.shape)
        #---------------------
        print("-------decoder--------")
        x = nn.MaxUnpool2d(kernel_size=2, stride=2, padding=0)(x, pool_indices[4])
        x = self.decoder1(x)
        print("x.shape: ",x.shape)
        x = nn.MaxUnpool2d(kernel_size=2, stride=2, padding=0)(x, pool_indices[3])
        x = self.decoder2(x)
        print("x.shape: ",x.shape)
        x = nn.MaxUnpool2d(kernel_size=2, stride=2, padding=0)(x, pool_indices[2])
        x = self.decoder3(x)
        print("x.shape: ",x.shape)
        x = nn.MaxUnpool2d(kernel_size=2, stride=2, padding=0)(x, pool_indices[1])
        x = self.decoder4(x)
        print("x.shape: ",x.shape)
        x = nn.MaxUnpool2d(kernel_size=2, stride=2, padding=0)(x, pool_indices[0])
        x = self.decoder5(x)
        print("x.shape: ",x.shape)
        return x
    
    def _initialize_weights(self, *stages):
        for modules in stages:
            for module in modules.modules():
                if isinstance(module, nn.Conv2d):
                    nn.init.kaiming_normal_(module.weight)
                    if module.bias is not None:
                        module.bias.data.zero_()
                elif isinstance(module, nn.BatchNorm2d):
                    module.weight.data.fill_(1)
                    module.bias.data.zero_()

# Example usage
if __name__ == "__main__":
    model = SegNet(num_classes=21)  # For example, Cityscapes dataset has 21 classes
    input_tensor = torch.randn(1, 3, 320, 480)
    output = model(input_tensor)
    print(output.shape)
# 输出
x.shape:  torch.Size([1, 64, 160, 240])
x.shape:  torch.Size([1, 128, 80, 120])
x.shape:  torch.Size([1, 256, 40, 60])
x.shape:  torch.Size([1, 512, 20, 30])
x.shape:  torch.Size([1, 512, 10, 15])
-------decoder--------
x.shape:  torch.Size([1, 512, 20, 30])
x.shape:  torch.Size([1, 256, 40, 60])
x.shape:  torch.Size([1, 128, 80, 120])
x.shape:  torch.Size([1, 64, 160, 240])
x.shape:  torch.Size([1, 21, 320, 480])
torch.Size([1, 21, 320, 480])

可以看到整个数据在编码器和解码器中的数据流转过程，最终输出为分为21类的结果。实际应用中，由于从头开始训练需要花不少时间，我们可以加载VGG模型的预训练权重，因为SegNet的编码器结构和VGG基本类似，可以稍作改动把五个编码层的权重（除最大池化层）替换为VGG的预训练权重。核心代码如下：

python 复制代码 
if self.preTrained:
    vgg = models.vgg16(pretrained=True)
else:
    vgg = models.vgg16(pretrained=False)
self.encoder1 = nn.Sequential(vgg.features[0:3])
self.encoder2 = nn.Sequential(vgg.features[5:8])
self.encoder3 = nn.Sequential(vgg.features[10:15])
self.encoder4 = nn.Sequential(vgg.features[17:22])
self.encoder5 = nn.Sequential(vgg.features[24:29])

下面我们看看SegNet的训练结果。

在VOC2012数据集上，SegNet和FCN都训练150个epoch，SegNet的效果是不如FCN的，可能是SegNet需要更多的资源，更长的训练轮数。

可以看到，在GID遥感数据集上，SegNet的分割效果就好了不少，同样的训练轮数和FCN效果类似，并且SegNet的边缘更平滑些。

相关推荐
咸鱼鲸21 分钟前
【PyTorch】PyTorch中数据准备工作（AI生成）
人工智能·pytorch·python
停走的风27 分钟前
二刷（李宏毅深度学习，醍醐灌顶，长刷长爽）
人工智能·深度学习
qinyia33 分钟前
Wisdom SSH：探索AI助手在复杂运维任务中的卓越表现
运维·人工智能·ssh
TY-202535 分钟前
二、深度学习——损失函数
人工智能·深度学习
Python×CATIA工业智造36 分钟前
列表页与详情页的智能识别：多维度判定方法与工业级实现
爬虫·深度学习·pycharm
京东零售技术42 分钟前
让大模型更懂你，京东零售的算法工程师做了这些事
人工智能·求职
PyAIExplorer43 分钟前
图像梯度处理与边缘检测：OpenCV 实战指南
人工智能·opencv·计算机视觉
CoovallyAIHub1 小时前
单目深度估计重大突破：无需标签，精度超越 SOTA！西湖大学团队提出多教师蒸馏新方案
深度学习·算法·计算机视觉
biubiubiu07061 小时前
微软云语音识别ASR示例Demo
人工智能·语音识别
CoovallyAIHub1 小时前
从FCOS3D到PGD：看深度估计如何快速搭建你的3D检测项目
深度学习·算法·计算机视觉
热门推荐
01深度神经网络训练过程与常见概念02KGG转MP3工具|非KGM文件|解密音频03Coze扣子平台完整体验和实践（附国内和国际版对比）04使用Ruby接入实时行情API教程05集群聊天服务器---MySQL数据库的建立06扣子（coze）实战|我用扣子搭建了一个自动分析小红薯笔记内容的AI应用|详细步骤拆解07Ubuntu环境下统计进程资源利用率08Java学习第十五部分——MyBatis09DeepSeek各版本说明与优缺点分析10【无标题】