Unet-語義分割

1.何為語義分割？

語義分割結合了目標檢測、圖像分類和圖像分割等技術。圖片輸入，通過語義分割模型對原有圖像分割成具有一定語義含義的區域塊，識別出每個區域塊語義類別，最終得到與原圖像等大小具有逐像素語義標注的分割圖像。

四幅圖分別代表(a)目標分類，(b)識別與定位，(c)語義分割，(d)實例分割

語義分割實則就是把整張圖片中的像素點分類。

例

原圖像                              標簽                

             

由圖像可知，共分為了四類:1.花瓣(紅色)，2.葉子(黃色)，3.莖稈(綠色)，4.背景(黑色)

一般情況下，數據集圖片大小可能不一，但導入網絡的數據集需要大小相同。此時需要對數據集進行圖片切割預處理。將數據集圖片等大小切割。

(跑一邊數據集，找到其中其中最小的W,H，選取切割大小 w<=W，h<=H)

 

 2.Unet

unet對稱語義分割模型

 

 主要由三部分組成。

第一部分主干特征提取(即灰線)，利用主干部分獲得一個又一個特征層，其主干特征提取部分與VGG16類似，利用此部分共獲得5個初步有效特征層，用於下一步與上采樣的的特征融合。

第二部分加強特征提取,將第一步獲得的5個特征層進行上采樣，並進行特征融合(即在channel通道維度進行堆疊)，獲得一個最終的，融合所有特征的有效特征層。

第三部分預測部分，通過獲得的最終有效特征層對每一個點進行分類(即對每個像素點分類)。

 

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import logging


class encoder(nn.Module):   # 下采樣部分
    def __init__(self, in_channels, out_channels):
        super(encoder, self).__init__()
        self.down_conv = nn.Sequential(     # 兩層卷積
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),   # 歸一化
            nn.ReLU(inplace=True),          # 激活函數
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
        # ceil_mode參數取整的時候向上取整，該參數默認為False表示取整的時候向下取整
        self.pool = nn.MaxPool2d(kernel_size=2, ceil_mode=True)

    def forward(self, x):
        out = self.down_conv(x)
        out_pool = self.pool(out)
        return out, out_pool   # 返回兩個 1.兩次池化后的，用於與后續上采樣相加
                                # 2.池化后的，用於繼續做下采樣


class decoder(nn.Module):   #  上采樣部分
    def __init__(self, in_channels, out_channels):
        super(decoder, self).__init__()
        # 反卷積
        self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)  # up-conv

        self.up_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x_copy, x, interpolate=True):
        out = self.up(x)
        if interpolate:
            # 迭代代替填充， 取得更好的結果
            out = F.interpolate(out, size=(x_copy.size(2), x_copy.size(3)),
                                mode="bilinear", align_corners=True
                                )
        else:
            # 如果填充物體積大小不同
            diffY = x_copy.size()[2] - x.size()[2]
            diffX = x_copy.size()[3] - x.size()[3]
            out = F.pad(out, (diffX // 2, diffX - diffX // 2, diffY, diffY - diffY // 2))
        # 連接
        out = torch.cat([x_copy, out], dim=1)  # 連接 按通道數疊加相當於羅列，非拼接
        out_conv = self.up_conv(out)  # 進行兩次卷積
        return out_conv


class BaseModel(nn.Module):
    def __init__(self):
        super(BaseModel, self).__init__()
        self.logger = logging.getLogger(self.__class__.__name__)

    def forward(self):
        raise NotImplementedError

    def summary(self):
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        nbr_params = sum([np.prod(p.size()) for p in model_parameters])
        self.logger.info(f'Nbr of trainable parametersL {nbr_params}')

    def __str__(self):
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        nbr_params = sum([np.prod(p.size()) for p in model_parameters])
        return super(BaseModel, self).__str__() + f"\nNbr of trainable parameters: {nbr_params}"


class UNet(BaseModel):
    def __init__(self, num_classes, in_channels=3, freeze_bn=False, **_):
        super(UNet, self).__init__()
        self.down1 = encoder(in_channels, 64)
        self.down2 = encoder(64, 128)
        self.down3 = encoder(128, 256)
        self.down4 = encoder(256, 512)
        self.middle_conv = nn.Sequential(
            nn.Conv2d(512, 1024, kernel_size=3, padding=1),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=True),
            nn.Conv2d(1024, 1024, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

        self.up1 = decoder(1024, 512)
        self.up2 = decoder(512, 256)
        self.up3 = decoder(256, 128)
        self.up4 = decoder(128, 64)
        self.final_conv = nn.Conv2d(64, num_classes, kernel_size=1)
        self._initalize_weights()
        if freeze_bn:
            self.freeze_bn()

    def _initalize_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
                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_()

    def forward(self, x):
        x1, x = self.down1(x)
        x2, x = self.down2(x)
        x3, x = self.down3(x)
        x4, x = self.down4(x)
        x = self.middle_conv(x)
        x = self.up1(x4, x)
        x = self.up2(x3, x)
        x = self.up3(x2, x)
        x = self.up4(x1, x)
        x = self.final_conv(x)
        return x

    def get_backbone_params(self):
        # There is no backbone for unet, all the parameters are trained from scratch
        return []

    def get_decoder_params(self):
        return self.parameters()

    def freeze_bn(self):
        for module in self.modules():
            if isinstance(module, nn.BatchNorm2d):
                module.eval()