• UNet语义分割网络

    1、本文参考

    [炼丹术]UNet图像分割模型相关总结_animalslin的技术博客_51CTO博客

    https://cuijiahua.com/blog/2019/11/dl-14.html

    Pytorch 深度学习实战教程（二）：UNet语义分割网络 - 腾讯云开发者社区-腾讯云

    2、UNet网络介绍

    UNet网络用于语义分割。

    语义就是给图像上目标类别中的每一点打一个标签，使得不同种类的东西在图像上被区分开来。可以理解成像素级别的分类任务，即对每个像素点进行分类。

    假如存在五类：Person（人）、Purse（包）、Plants/Grass（植物/草）、Sidewalk（人行道）、Building/Structures（建筑物）。需要创建一个one-hot编码的目标类别标注，即为每个类别创建一个输出通道。因为有5个类别，所以网络输出的通道数也为5，如下图所示：

     

     

    因为不存在同一个像素点在两个以上的通道均为1的情况（存疑），所以预测的结果可以通过对每个像素在深度上求argmax的方式被整合到一张分割图中，进而可以通过重叠的方式观察到每个目标。

    UNet网络的架构如下（实际实施时思想不变，但是略有调整）：

     

    3、UNet训练整体方案

    （1）通过labelme进行语义标注，产出结果json文件

    （2）编写代码，根据json文件的points信息，从原图中获取mask图片

    （3）在UNet网络中，输入3通道图片，输出预测的1通道mask（假定只有一个识别类别），将预测的mask和实际的mask计算BCELoss从而进行拟合操作，并且输出准确率和dice score的监控指标

    4、UNet网络实施分析

    （1）labelme进行多边形标注

     标注完成后，会在图片所在目录下生成json文件。

    （2）根据json文件生成mask图片

    文件名：json2mask.py

    1. import os
    2. import cv2
    3. import numpy as np
    4. from PIL import Image, ImageDraw
    5. import json
    6.  
    7. CLASS_NAMES = ['dog', 'cat']
    8.  
    9. def make_mask(image_dir, save_dir):
    10.     data = os.listdir(image_dir)
    11.     temp_data = []
    12.     for i in data:
    13.         if i.split('.')[1] == 'json':
    14.             temp_data.append(i)
    15.         else:
    16.             continue
    17.     for js in temp_data:
    18.         json_data = json.load(open(os.path.join(image_dir, js), 'r'))
    19.         shapes_ = json_data['shapes']
    20.         mask = Image.new('P', Image.open(os.path.join(image_dir, js.replace('json', 'jpg'))).size)
    21.         for shape_ in shapes_:
    22.             label = shape_['label']
    23.             points = shape_['points']
    24.             points = tuple(tuple(i) for i in points)
    25.             mask_draw = ImageDraw.Draw(mask)  # 类似于函数声明
    26.             mask_draw.polygon(points, fill=CLASS_NAMES.index(label) + 1)
    27.             mask = np.array(mask) * 255
    28.             cv2.imshow('mask', mask)
    29.             cv2.waitKey(0)
    30.         cv2.imwrite(os.path.join(save_dir, js.replace('json', 'jpg')), mask)
    31.  
    32. def vis_label(img):
    33.     img = Image.open(img)
    34.     img = np.array(img)
    35.     print(set(img.reshape(-1).tolist()))
    36.  
    37. if __name__ == '__main__':
    38.     make_mask('D:\\ai_data\\cat\\val', 'D:\\ai_data\\cat\\val_mask')

    说明：

    • Image.new中mode='P'，代表生成的图片为8-bit pixels，适合用于生成mask图片
    • 像素值=255表示白色，也就是说mask图片中mask部分为白色，非mask部分为黑色。实际得到的mask图片中mask会存在(249,255)的值，使用时需要再处理下。

    （3）UNet网络构造

    1. import torch
    2. import torch.nn as nn
    3. import torch.nn.functional as F
    4.  
    5. class DoubleConv(nn.Module):
    6.     """(convolution => [BN] => ReLU) * 2"""
    7.  
    8.     def __init__(self, in_channels, out_channels):
    9.         super().__init__()
    10.         self.double_conv = nn.Sequential(
    11.             nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0),
    12.             nn.BatchNorm2d(out_channels),
    13.             nn.ReLU(inplace=True),
    14.             nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0),
    15.             nn.BatchNorm2d(out_channels),
    16.             nn.ReLU(inplace=True)
    17.         )
    18.  
    19.     def forward(self, x):
    20.         return self.double_conv(x)
    21.  
    22. class Down(nn.Module):
    23.     """Downscaling with maxpool then double conv"""
    24.  
    25.     def __init__(self, in_channels, out_channels):
    26.         super().__init__()
    27.         self.maxpool_conv = nn.Sequential(
    28.             nn.MaxPool2d(2),
    29.             DoubleConv(in_channels, out_channels)
    30.         )
    31.  
    32.     def forward(self, x):
    33.         return self.maxpool_conv(x)
    34.  
    35.  
    36. class Up(nn.Module):
    37.     """Upscaling then double conv"""
    38.  
    39.     def __init__(self, in_channels, out_channels, bilinear=True):
    40.         super().__init__()
    41.  
    42.         if bilinear:
    43.             self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
    44.         else:
    45.             self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
    46.  
    47.         self.conv = DoubleConv(in_channels, out_channels)
    48.  
    49.     def forward(self, x1, x2):
    50.         x1 = self.up(x1)
    51.         # input is NCHW
    52.         diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
    53.         diffX = torch.tensor([x2.size()[3] - x1.size()[3]])
    54.  
    55.         x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
    56.                         diffY // 2, diffY - diffY // 2])
    57.         x = torch.cat([x2, x1], dim=1)
    58.         return self.conv(x)
    59.  
    60. class OutConv(nn.Module):
    61.     def __init__(self, in_channels, out_channels):
    62.         super(OutConv, self).__init__()
    63.         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
    64.  
    65.     def forward(self, x):
    66.         return self.conv(x)
    67.  
    68. class UNet(nn.Module):
    69.  
    70.     def __init__(self, n_channels, n_classes, bilinear=False):
    71.         super(UNet, self).__init__()
    72.         self.n_channels = n_channels
    73.         self.n_classes = n_classes
    74.         self.bilinear = bilinear
    75.  
    76.         self.inc = DoubleConv(n_channels, 64)
    77.         self.down1 = Down(64, 128)
    78.         self.down2 = Down(128, 256)
    79.         self.down3 = Down(256, 512)
    80.         self.down4 = Down(512, 1024)
    81.         self.up1 = Up(1024, 512, bilinear)
    82.         self.up2 = Up(512, 256, bilinear)
    83.         self.up3 = Up(256, 128, bilinear)
    84.         self.up4 = Up(128, 64, bilinear)
    85.         self.outc = OutConv(64, n_classes)
    86.  
    87.     def forward(self, x):
    88.         x1 = self.inc(x)
    89.         x2 = self.down1(x1)
    90.         x3 = self.down2(x2)
    91.         x4 = self.down3(x3)
    92.         x5 = self.down4(x4)
    93.         x = self.up1(x5, x4)
    94.         x = self.up2(x, x3)
    95.         x = self.up3(x, x2)
    96.         x = self.up4(x, x1)
    97.         logits = self.outc(x)
    98.         return logits
    99.  
    100. if __name__ == '__main__':
    101.     net = UNet(n_channels=3, n_classes=1)
    102.     print(net)
    103.  
    104.     x = torch.randn([1, 3, 572, 572])
    105.     out = net(x)
    106.     print(out.shape)

     说明：

    • 本代码为按照UNet论文构造的网络，实际中并未使用该代码，需要稍作修改
    • 本代码适合阅读UNet网络，不明白之处可参考：Pytorch 深度学习实战教程（二）：UNet语义分割网络 - 腾讯云开发者社区-腾讯云
    •  下采样通过卷积核最大池化完成
    • 上采样通过转置卷积以及和下采样的特征concat完成。 
    • 上采样是会将通道维度减少一半，比如1024到512，因为和下采样的特征(同样也是512)在dim=1(channel维度)进行了concat，所以channel维度的值又变为了1024.
    • 论文的输入w、h和输出的w‘、h’大小不一样，在mask图片比对时会有问题，所以我们希望输入和输出的wh保持一致。此时会设置padding=1，这样double_conv时候w、h会保持不变，只有池化时变为原来的一半，上采样时候又会变为原来的两倍。
    • 以上代码只作为了解UNet网络使用，不作为整个工程的代码。

    （4）主函数train.py

    1. import torch
    2. import albumentations as A
    3. from albumentations.pytorch import ToTensorV2
    4. from tqdm import tqdm
    5. import torch.nn as nn
    6. import torch.optim as optim
    7. from model import UNET
    8. # from unet_model_new import UNet
    9. from utils import (
    10.     load_checkpoint,
    11.     save_checkpoint,
    12.     get_loaders,
    13.     check_accuracy,
    14.     save_predictions_as_imgs,
    15. )
    16.  
    17. # 超参
    18. learning_rate = 1e-4
    19. device = 'cpu'
    20. batch_size = 1
    21. num_epochs = 30
    22. num_workers = 0
    23. image_height = 160
    24. image_width = 240
    25. pin_memory = False
    26. load_model = False
    27. train_img_dir = "D:\\ai_data\\cat\\train2"
    28. train_mask_dir = "D:\\ai_data\\cat\\train2_mask"
    29. val_img_dir = "D:\\ai_data\\cat\\val2"
    30. val_mask_dir = "D:\\ai_data\\cat\\val2_mask"
    31.  
    32.  
    33. def train_fn(loader, model, optimizer, loss_fn):
    34.  
    35.     for batch_idx, (data, targets) in enumerate(tqdm(loader)):
    36.         data = data.to(device=device)
    37.         targets = targets.float().unsqueeze(1).to(device=device)
    38.  
    39.         predictions = model(data)
    40.         loss = loss_fn(predictions, targets)
    41.  
    42.         optimizer.zero_grad()
    43.         loss.backward()
    44.  
    45.  
    46. def main():
    47.     train_transform = A.Compose(
    48.         [
    49.             A.Resize(height=image_height, width=image_width),
    50.             A.Rotate(limit=35, p=1.0),
    51.             A.HorizontalFlip(p=0.5),
    52.             A.VerticalFlip(p=0.1),
    53.             A.Normalize(
    54.                 mean=[0.0, 0.0, 0.0],
    55.                 std=[1.0, 1.0, 1.0],
    56.                 max_pixel_value=255.0
    57.             ),
    58.             ToTensorV2(),
    59.         ],
    60.     )
    61.  
    62.     val_transform = A.Compose(
    63.         [
    64.             A.Resize(height=image_height, width=image_width),
    65.             A.Normalize(
    66.                 mean=[0.0, 0.0, 0.0],
    67.                 std=[1.0, 1.0, 1.0],
    68.                 max_pixel_value=255.0
    69.             ),
    70.             ToTensorV2(),
    71.         ],
    72.     )
    73.  
    74.     model = UNET(in_channels=3, out_channels=1).to(device)
    75.     loss_fn = nn.BCEWithLogitsLoss()
    76.     optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    77.  
    78.     train_loader, val_loader = get_loaders(
    79.         train_img_dir,
    80.         train_mask_dir,
    81.         val_img_dir,
    82.         val_mask_dir,
    83.         batch_size,
    84.         train_transform,
    85.         val_transform,
    86.         num_workers,
    87.         pin_memory
    88.     )
    89.  
    90.     if load_model:
    91.         load_checkpoint(torch.load("my_checkpoint.pth.tar"), model)
    92.  
    93.     check_accuracy(-1, "val", val_loader, model, device=device)
    94.  
    95.     for epoch in range(num_epochs):
    96.         train_fn(train_loader, model, optimizer, loss_fn)
    97.  
    98.         checkpoint = {
    99.             "state_dict": model.state_dict(),
    100.             "optimizer": optimizer.state_dict(),
    101.         }
    102.         save_checkpoint(checkpoint)
    103.  
    104.         check_accuracy(epoch, "train", train_loader, model, device=device)
    105.         check_accuracy(epoch, "val", val_loader, model, device=device)
    106.  
    107.         save_predictions_as_imgs(val_loader, model, folder="saved_images/", device=device)
    108.  
    109.  
    110. if __name__ == "__main__":
    111.     main()

     （5）数据加载dataset.py

    1. import os
    2. from PIL import Image
    3. from torch.utils.data import Dataset
    4. import numpy as np
    5.  
    6. class CarvanaDataset(Dataset):
    7.     def __init__(self, image_dir, mask_dir, transform=None):
    8.         self.image_dir = image_dir
    9.         self.mask_dir = mask_dir
    10.         self.transform = transform
    11.         self.images = os.listdir(image_dir)
    12.  
    13.     def __len__(self):
    14.         return len(self.images)
    15.  
    16.     def __getitem__(self, index):
    17.         img_path = os.path.join(self.image_dir, self.images[index])
    18.         mask_path = os.path.join(self.mask_dir, self.images[index].replace(".jpg", "_mask.jpg"))
    19.         image = np.array(Image.open(img_path).convert("RGB"))
    20.         mask = np.array(Image.open(mask_path).convert("L"), dtype=np.float32)
    21.         mask[mask > 200.0] = 1.0   # 转换为灰度图后并非全是255白色
    22.  
    23.         if self.transform is not None:
    24.             augmentations = self.transform(image=image, mask=mask)
    25.             image = augmentations["image"]
    26.             mask = augmentations["mask"]
    27.  
    28.         return image, mask

    （6）模型model.py

    1. import torch
    2. import torch.nn as nn
    3. import torch.functional as F
    4. import torchvision.transforms.functional as TF
    5.  
    6. class DoubleConv(nn.Module):
    7.     def __init__(self, in_channels, out_channels):
    8.         super(DoubleConv, self).__init__()
    9.         self.conv = nn.Sequential(
    10.             nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False),    # padding=1,保证conv2d的输出hw保持不变
    11.             nn.BatchNorm2d(out_channels),
    12.             nn.ReLU(inplace=True),
    13.             nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False),   # padding=1,保证conv2d的输出hw保持不变
    14.             nn.BatchNorm2d(out_channels),
    15.             nn.ReLU(inplace=True),
    16.         )
    17.  
    18.     def forward(self, x):
    19.         return self.conv(x)
    20.  
    21. class UNET(nn.Module):
    22.     def __init__(
    23.             self, in_channels=3, out_channels=1, features=[64, 128, 256, 512],
    24.     ):
    25.         super(UNET, self).__init__()
    26.         self.ups = nn.ModuleList()
    27.         self.downs = nn.ModuleList()
    28.         self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
    29.  
    30.         # Down part of UNET
    31.         for feature in features:
    32.             self.downs.append(DoubleConv(in_channels, feature))
    33.             in_channels = feature
    34.  
    35.         # Up part of UNET
    36.         for feature in reversed(features):
    37.             self.ups.append(
    38.                 nn.ConvTranspose2d(
    39.                     feature*2, feature, kernel_size=2, stride=2,
    40.                 )
    41.             )
    42.             self.ups.append(DoubleConv(feature*2, feature))
    43.  
    44.         self.bottleneck = DoubleConv(features[-1], features[-1]*2)
    45.         self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
    46.  
    47.     def forward(self, x):
    48.         skip_connections = []
    49.  
    50.         for down in self.downs:
    51.             x = down(x)
    52.             skip_connections.append(x)
    53.             x = self.pool(x)
    54.  
    55.         x = self.bottleneck(x)
    56.         skip_connections = skip_connections[::-1]
    57.  
    58.         for idx in range(0, len(self.ups), 2):
    59.             x = self.ups[idx](x)
    60.             skip_connection = skip_connections[idx//2]
    61.  
    62.             if x.shape != skip_connection.shape:
    63.                 x = TF.resize(x, size=skip_connection.shape[2:])  # 因为有padding=1，所以到不了这一步
    64.                 # diffY = torch.tensor([skip_connection.size()[2] - x.size()[2]])
    65.                 # diffX = torch.tensor([skip_connection.size()[3] - x.size()[3]])
    66.                 # x = F.pad(x, [diffX // 2, diffX - diffX // 2,
    67.                 #         diffY // 2, diffY - diffY // 2])
    68.  
    69.             concat_skip = torch.cat((skip_connection, x), dim=1)
    70.             x = self.ups[idx+1](concat_skip)
    71.  
    72.         return self.final_conv(x)
    73.  
    74. def test():
    75.     x = torch.randn((3, 1, 572, 572))
    76.     model = UNET(in_channels=1, out_channels=1)
    77.     preds = model(x)
    78.     assert preds.shape == x.shape
    79.  
    80. if __name__ == "__main__":
    81.     test()

    （7）工具utils.py

    1. import torch
    2. import torchvision
    3. from dataset import CarvanaDataset
    4. from torch.utils.data import DataLoader
    5.  
    6. def save_checkpoint(state, filename="my_checkpoint.pth.tar"):
    7.     print("=> Saving checkpoint")
    8.     torch.save(state, filename)
    9.  
    10. def load_checkpoint(checkpoint, model):
    11.     print("=> Loading checkpoint")
    12.     model.load_state_dict(checkpoint["state_dict"])
    13.  
    14. def get_loaders(
    15.     train_dir,
    16.     train_maskdir,
    17.     val_dir,
    18.     val_maskdir,
    19.     batch_size,
    20.     train_transform,
    21.     val_transform,
    22.     num_workers=4,
    23.     pin_memory=True,
    24. ):
    25.     train_ds = CarvanaDataset(
    26.         image_dir=train_dir,
    27.         mask_dir=train_maskdir,
    28.         transform=train_transform,
    29.     )
    30.  
    31.     train_loader = DataLoader(
    32.         train_ds,
    33.         batch_size=batch_size,
    34.         num_workers=num_workers,
    35.         pin_memory=pin_memory,
    36.         shuffle=True,
    37.     )
    38.  
    39.     val_ds = CarvanaDataset(
    40.         image_dir=val_dir,
    41.         mask_dir=val_maskdir,
    42.         transform=val_transform,
    43.     )
    44.  
    45.     val_loader = DataLoader(
    46.         val_ds,
    47.         batch_size=batch_size,
    48.         num_workers=num_workers,
    49.         pin_memory=pin_memory,
    50.         shuffle=False,
    51.     )
    52.  
    53.     return train_loader, val_loader
    54.  
    55. def check_accuracy(epoch, attr, loader, model, device="cuda"):
    56.     num_correct = 0
    57.     num_pixels = 0
    58.     dice_score = 0
    59.     model.eval()
    60.  
    61.     with torch.no_grad():
    62.         for x, y in loader:
    63.             x = x.to(device)
    64.             y = y.to(device).unsqueeze(1)
    65.             preds = torch.sigmoid(model(x))
    66.             preds = (preds > 0.5).float()
    67.             num_correct += (preds == y).sum()
    68.             num_pixels += torch.numel(preds)
    69.             dice_score += (2 * (preds * y).sum()) / (
    70.                 (preds + y).sum() + 1e-8
    71.             )
    72.  
    73.     print(f"{attr}_{epoch+1}: Got {num_correct}/{num_pixels} with acc {num_correct/num_pixels*100:.2f}")
    74.     print(f"{attr}_{epoch+1}: Dice score: {dice_score/len(loader)}")
    75.     model.train()
    76.  
    77. def save_predictions_as_imgs(
    78.     loader, model, folder="saved_images/", device="cuda"
    79. ):
    80.     model.eval()
    81.     for idx, (x, y) in enumerate(loader):
    82.         x = x.to(device=device)
    83.         with torch.no_grad():
    84.             preds = torch.sigmoid(model(x))
    85.             preds = (preds > 0.5).float()
    86.         torchvision.utils.save_image(
    87.             preds, f"{folder}/pred_{idx}.png"
    88.         )
    89.         torchvision.utils.save_image(y.unsqueeze(1), f"{folder}{idx}.png")
    90.  
    91.     model.train()

    （8）监控指标dice score说明

    参考文档：关于图像分割的评价指标dice_Pierce_KK的博客-CSDN博客_dice评价指标

     dice指标也用在机器学习中，它的表达式为：

    这与机器学习中的评价指标F1是相同的。

    准确率指标：

    召回率指标：

     

    F1则是基于准确率和召回率的调和平均值，即：

     

     

    dice指标是医学图像中的常见指标，常用于评价图像分割算法的好坏。从公式上来做直观的理解，如下图所示，其代表的是两个体相交的面积占总面积的比值，完美分割该值为1.

     

     本试验中，准确率能够达到60%+，disc score只有0.4+，整体效果不佳。

     

    5、 UNet后续发展

    （1）UNet网络的思想：

    • 下采样+上采样作为整体的网络结构（Encoder-Decoder）
    • 多尺度的特征融合
    • 信息流通的方式
    • 获得像素级别的segment map

    （2）对于改进UNet的见解，参考：谈一谈UNet图像分割_3Ｄ视觉工坊的博客-CSDN博客

    很多人都喜欢在UNet进行改进，换个优秀的编码器，然后自己在手动把解码器对应实现一下。执御为什么选择UNet上进行改进，可能是因为UNet网络的结构比较简单，而且UNet的效果在很多场景下的表现可能都是差强人意的。 

     UNet最原始的设计思路，相对于后面系列的一个劣势就是：信息融合、位置不偏移。

  • 相关阅读:
    【CTF Web】CTFShow web18 Writeup（文件包含漏洞+日志注入+RCE）
    125. 验证回文串
    ubuntu18.04下zookeeper安装与简单使用
    Ubuntu20.04 从头到尾完整版安装anaconda、cuda、cudnn、pytorch、paddle2.3成功记录
    JS 方法实现复制粘贴
    【lambda表达式】常用的函数式接口
    【csdn】gitcode初体验（开发云、Pages等）（持续更新）
    Nacos-Feign-Gateway
    【Python爬虫原理与基本请求库urllib详解】
    不用任何 js 库，纯前端导出数据到 Excel / CSV 文件就这么简单几行代码
  • 原文地址：https://blog.csdn.net/benben044/article/details/126142576