语义分割 - 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以上的期刊论文进行复现。非常适合初学者,或者急需写文章却没有头绪的小伙伴。订阅价格很便宜,一顿午饭钱,就可以帮你节省数周甚至数月时间。希望大家多多关注,相互讨论,相互进步。