devbox
我家小花儿
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

UNet语义分割网络

作者：我家小花儿 | 2024-02-15 22:31:42

赞

踩

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最原始的设计思路，相对于后面系列的一个劣势就是：信息融合、位置不偏移。

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/我家小花儿/article/detail/87672?site
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号