语义分割基础篇系列 — Unet（一）

语义分割 - Unet（一）

---

---

前言

本篇文章主要介绍Unet模型的基础概念以及Unet项目的实现，旨在帮助初学者快速上手。Unet网络是最经典的编解码结构，左侧部分负责特征提取，右侧部分负责特征图的上采样，简单有效的结构是语义分割必掌握的结构之一。在它问世十年以来，无论是工业界还是学术界对它的优化和创新从未停止，一直到现在该结构及其衍生体仍然是最活跃的分割模型。

原文链接：https://arxiv.org/abs/1505.04597
代码链接：https://github.com/bubbliiiing/unet-pytorch.git

---

一、概述

1.1摘要

普遍认为，成功训练深度网络需要成千上万个带有注释的训练样本。在本文中，我们提出了一种网络和训练策略，依赖于强大的数据增强，以更有效地利用可用的带有注释的样本。该架构包括一个用于捕获上下文的收缩路径和一个对称的扩展路径，可实现精确定位。我们展示了这样的网络可以从极少量的图像进行端到端训练，并在ISBI挑战中超过先前最佳方法（滑动窗口卷积网络）用于电子显微镜堆叠中神经结构的分割。使用相同网络在透射光显微镜图像（相差和差示显微镜）上训练，我们在这些类别中以较大的优势赢得了2015年的ISBI细胞跟踪挑战。此外，该网络速度很快。对于一个512x512的图像，使用最新的GPU只需不到一秒钟进行分割。

1.2结构简介

在Unet中，网络对输入图像进行4次降采样，得到图像的高级特征，这个阶段可以替换为任何骨干网络，如VGG、ResNet、MobileNet等。然后通过上采样来逐步恢复图像的细节，并通过skip connect(图中的灰色箭头)的方式来连接编码过程中的相应尺度的特征图，这种做法可以有效地增加特征图的信息。

二、使用步骤

1.下载项目包

使用git或者download zip都行。

2.自定义数据集

(1) 安装labelme，选择Edit–>Create Polygons，然后进行标注

pip install labelme==3.16.7
labelme

(2)自定义数据集
将标注好的数据放在 datasets/before
然后修改json_to_dataset.py第23行，将classes中修改成自己的标签名称，第一个类别__background__不变，后面依次修改成自己的标签，再运行该文件。

if __name__ == '__main__':
    jpgs_path   = "datasets/JPEGImages"
    pngs_path   = "datasets/SegmentationClass"
    # 修改classes 
    classes     = ["_background_","person","computer"]

然后datasets/JPEGImages和datasets/SegmentationClass中会生成对应的原图和label文件。
然后生成对应的train.txt和val.txt

import os
import random 
imgdir="datasets/SegmentationClass"
imglist=os.listdir(imgdir)
imglist = [imgname[:-4]+'\n' for imgname in imglist]
train_sample = random.sample(imglist, int(0.8* len(imglist)))
with open("datasets/train.txt", 'w') as f:
    f.writelines(train_sample)
valid_sample = []
for name in imglist:
    if name not in train_sample:
        valid_sample.append(name)
with open("datasets/val.txt", 'w') as f:
    f.writelines(valid_sample)

最后数据集格式为

3.模型训练

训练时需要注意以下几点
(1)num_classes = 分割的类别+1，因为0为背景
(2)该项目并不是原本Unet模型，其backbone为vgg或者resnet50，如果需要model_path，则需要去github上下载对应的权重
(3)训练时需要修改以下几处数据集路径。
train.py中

    #   数据集路径
    #------------------------------#
    VOCdevkit_path  = 'datasets'

    #   读取数据集对应的txt
    #---------------------------#
    with open(os.path.join(VOCdevkit_path, "train.txt"),"r") as f:
        train_lines = f.readlines()
    with open(os.path.join(VOCdevkit_path, "val.txt"),"r") as f:
        val_lines = f.readlines()

utils/dataloader.py中

        #   从文件中读取图像
        #-------------------------------#
        jpg         = Image.open(os.path.join(os.path.join(self.dataset_path, "JPEGImages"), name + ".jpg"))
        png         = Image.open(os.path.join(os.path.join(self.dataset_path, "SegmentationClass"), name + ".png"))

还有utils/callbacks中

    def on_epoch_end(self, epoch, model_eval):
        if epoch % self.period == 0 and self.eval_flag:
            self.net    = model_eval
            gt_dir      = os.path.join(self.dataset_path, "SegmentationClass/")  #此处需要注意
            pred_dir    = os.path.join(self.miou_out_path, 'detection-results')
            if not os.path.exists(self.miou_out_path):
                os.makedirs(self.miou_out_path)
            if not os.path.exists(pred_dir):
                os.makedirs(pred_dir)
            print("Get miou.")
            for image_id in tqdm(self.image_ids):
                #-------------------------------#
                #   从文件中读取图像
                #-------------------------------#
                image_path  = os.path.join(self.dataset_path, "JPEGImages/"+image_id+".jpg") #此处需要注意
                image       = Image.open(image_path)
                #------------------------------#
                #   获得预测txt
                #------------------------------#
                image       = self.get_miou_png(image)
                image.save(os.path.join(pred_dir, image_id + ".png"))

然后就可以进行训练

4.模型推理

模型推理也非常简单，主要是涉及到unet.py和predict.py，代码中也有详细的注释，按照需求修改即可。

三、后续内容

有改进需求或者发论文需求的同学，可订阅本人的《手把手复现语义分割算法》专栏。该专栏主要是针对知网上核心级别，影响因子1.0以上的期刊论文进行复现。非常适合初学者，或者急需写文章却没有头绪的小伙伴。订阅价格很便宜，一顿午饭钱，就可以帮你节省数周甚至数月时间。希望大家多多关注，相互讨论，相互进步。