全卷积网络(FCN)实战:使用FCN实现语义分割

虚幻大学 xuhss 515℃ 0评论

Python微信订餐小程序课程视频

https://edu.csdn.net/course/detail/36074

Python实战量化交易理财系统

https://edu.csdn.net/course/detail/35475

摘要:FCN对图像进行像素级的分类,从而解决了语义级别的图像分割问题。

本文分享自华为云社区《全卷积网络(FCN)实战:使用FCN实现语义分割》,作者: AI浩。

FCN对图像进行像素级的分类,从而解决了语义级别的图像分割(semantic segmentation)问题。与经典的CNN在卷积层之后使用全连接层得到固定长度的特征向量进行分类(全联接层+softmax输出)不同,FCN可以接受任意尺寸的输入图像,采用反卷积层对最后一个卷积层的feature map进行上采样, 使它恢复到输入图像相同的尺寸,从而可以对每个像素都产生了一个预测, 同时保留了原始输入图像中的空间信息, 最后在上采样的特征图上进行逐像素分类。

下图是语义分割所采用的全卷积网络(FCN)的结构示意图:

5cc93af3970d33d8adc6a4a4f16ee8ea - 全卷积网络(FCN)实战:使用FCN实现语义分割

传统的基于CNN的分割方法缺点?

传统的基于CNN的分割方法:为了对一个像素分类,使用该像素周围的一个图像块作为CNN的输入,用于训练与预测,这种方法主要有几个缺点:

1)存储开销大,例如,对每个像素使用15 * 15的图像块,然后不断滑动窗口,将图像块输入到CNN中进行类别判断,因此,需要的存储空间随滑动窗口的次数和大小急剧上升;

2)效率低下,相邻像素块基本上是重复的,针对每个像素块逐个计算卷积,这种计算有很大程度上的重复;

3)像素块的大小限制了感受区域的大小,通常像素块的大小比整幅图像的大小小很多,只能提取一些局部特征,从而导致分类性能受到限制。而全卷积网络(FCN)则是从抽象的特征中恢复出每个像素所属的类别。即从图像级别的分类进一步延伸到像素级别的分类。

FCN改变了什么?

对于一般的分类CNN网络,如VGG和Resnet,都会在网络的最后加入一些全连接层,经过softmax后就可以获得类别概率信息。但是这个概率信息是1维的,即只能标识整个图片的类别,不能标识每个像素点的类别,所以这种全连接方法不适用于图像分割。

而FCN提出可以把后面几个全连接都换成卷积,这样就可以获得一张2维的feature map,后接softmax层获得每个像素点的分类信息,从而解决了分割问题,如图。

a75a16746b31f5c7be8d074cdbd0dd01 - 全卷积网络(FCN)实战:使用FCN实现语义分割

FCN缺点

(1)得到的结果还是不够精细。进行8倍上采样虽然比32倍的效果好了很多,但是上采样的结果还是比较模糊和平滑,对图像中的细节不敏感。(2)对各个像素进行分类,没有充分考虑像素与像素之间的关系。忽略了在通常的基于像素分类的分割方法中使用的空间规整(spatial regularization)步骤,缺乏空间一致性。

数据集

本例的数据集采用PASCAL VOC 2012 数据集,它有二十个类别:

Person:person

Animal:bird, cat, cow, dog, horse, sheep

Vehicle:aeroplane, bicycle, boat, bus, car, motorbike, train

Indoor:bottle, chair, dining table, potted plant, sofa, tv/monitor

01442286d296a7ead2f21995f229c035 - 全卷积网络(FCN)实战:使用FCN实现语义分割

下载地址:The PASCAL Visual Object Classes Challenge 2012 (VOC2012) (ox.ac.uk)。

数据集的结构:

VOCdevkit
 └── VOC2012
 ├── Annotations 所有的图像标注信息(XML文件)
 ├── ImageSets 
 │ ├── Action 人的行为动作图像信息
 │ ├── Layout 人的各个部位图像信息
 │ │
 │ ├── Main 目标检测分类图像信息
 │ │ ├── train.txt 训练集(5717)
 │ │ ├── val.txt 验证集(5823)
 │ │ └── trainval.txt 训练集+验证集(11540)
 │ │
 │ └── Segmentation 目标分割图像信息
 │ ├── train.txt 训练集(1464)
 │ ├── val.txt 验证集(1449)
 │ └── trainval.txt 训练集+验证集(2913)
 │ 
 ├── JPEGImages 所有图像文件
 ├── SegmentationClass 语义分割png图(基于类别)
 └── SegmentationObject 实例分割png图(基于目标)

数据集包含物体检测和语义分割,我们只需要语义分割的数据集,所以可以考虑把多余的图片删除,删除的思路:

1、获取所有图片的name。

2、获取所有语义分割mask的name。

3、求二者的差集,然后将差集的name删除。

代码如下:

import glob
import os
image\_all = glob.glob('data/VOCdevkit/VOC2012/JPEGImages/*.jpg')
image\_all\_name = [image\_file.replace('\\', '/').split('/')[-1].split('.')[0] for image\_file in image\_all]

image\_SegmentationClass = glob.glob('data/VOCdevkit/VOC2012/SegmentationClass/*.png')
image\_se\_name= [image\_file.replace('\\', '/').split('/')[-1].split('.')[0] for image\_file in image\_SegmentationClass]
image\_other=list(set(image\_all\_name) - set(image\_se\_name))
print(image\_other)
for image\_name in image\_other:
 os.remove('data/VOCdevkit/VOC2012/JPEGImages/{}.jpg'.format(image\_name))

代码链接

本例选用的代码来自deep-learning-for-image-processing/pytorch_segmentation/fcn at master · WZMIAOMIAO/deep-learning-for-image-processing (github.com)

其他的代码也有很多,这篇比较好理解!

其实还有个比较好的图像分割库:https://github.com/qubvel/segmentation_models.pytorch

这个图像分割集合由俄罗斯的程序员小哥Pavel Yakubovskiy一手打造。在后面的文章,我也会使用这个库演示。

项目结构

├── src: 模型的backbone以及FCN的搭建
├── train\_utils: 训练、验证以及多GPU训练相关模块
├── my\_dataset.py: 自定义dataset用于读取VOC数据集
├── train.py: 以fcn\_resnet50(这里使用了Dilated/Atrous Convolution)进行训练
├── predict.py: 简易的预测脚本,使用训练好的权重进行预测测试
├── validation.py: 利用训练好的权重验证/测试数据的mIoU等指标,并生成record\_mAP.txt文件
└── pascal\_voc\_classes.json: pascal\_voc标签文件

由于代码很多不能一一讲解,所以,接下来对重要的代码做剖析。

自定义数据集读取

my_dataset.py自定义数据读取的方法,代码如下:

import os
import torch.utils.data as data
from PIL import Image

class VOCSegmentation(data.Dataset):
 def \_\_init\_\_(self, voc\_root, year="2012", transforms=None, txt\_name: str = "train.txt"):
 super(VOCSegmentation, self).\_\_init\_\_()
 assert year in ["2007", "2012"], "year must be in ['2007', '2012']"
 root = os.path.join(voc\_root, "VOCdevkit", f"VOC{year}")
 root=root.replace('\\','/')
 assert os.path.exists(root), "path '{}' does not exist.".format(root)
 image\_dir = os.path.join(root, 'JPEGImages')
 mask\_dir = os.path.join(root, 'SegmentationClass')

 txt\_path = os.path.join(root, "ImageSets", "Segmentation", txt\_name)
 txt\_path=txt\_path.replace('\\','/')
 assert os.path.exists(txt\_path), "file '{}' does not exist.".format(txt\_path)
 with open(os.path.join(txt\_path), "r") as f:
 file\_names = [x.strip() for x in f.readlines() if len(x.strip()) > 0]

 self.images = [os.path.join(image\_dir, x + ".jpg") for x in file\_names]
 self.masks = [os.path.join(mask\_dir, x + ".png") for x in file\_names]
 assert (len(self.images) == len(self.masks))
 self.transforms = transforms

导入需要的包。

定义VOC数据集读取类VOCSegmentation。在init方法中,核心是读取image列表和mask列表。

 def \_\_getitem\_\_(self, index):
 img = Image.open(self.images[index]).convert('RGB')
 target = Image.open(self.masks[index])

 if self.transforms is not None:
 img, target = self.transforms(img, target)
 return img, target

__getitem__方法是获取单张图片和图片对应的mask,然后对其做数据增强。

 def collate\_fn(batch):
 images, targets = list(zip(*batch))
 batched\_imgs = cat\_list(images, fill\_value=0)
 batched\_targets = cat\_list(targets, fill\_value=255)
 return batched\_imgs, batched\_targets

collate_fn方法是对一个batch中数据调用cat_list做数据对齐。

在train.py中torch.utils.data.DataLoader调用

 train\_loader = torch.utils.data.DataLoader(train\_dataset,
 batch\_size=batch\_size,
 num\_workers=num\_workers,
 shuffle=True,
 pin\_memory=True,
 collate\_fn=train\_dataset.collate\_fn)
 val\_loader = torch.utils.data.DataLoader(val\_dataset,
 batch\_size=1,
 num\_workers=num\_workers,
 pin\_memory=True,
 collate\_fn=val\_dataset.collate\_fn)

训练

重要参数

打开train.py,我们先认识一下重要的参数:

def parse\_args():
 import argparse
 parser = argparse.ArgumentParser(description="pytorch fcn training")
 # 数据集的根目录(VOCdevkit)所在的文件夹
 parser.add\_argument("--data-path", default="data/", help="VOCdevkit root")
 parser.add\_argument("--num-classes", default=20, type=int)
 parser.add\_argument("--aux", default=True, type=bool, help="auxilier loss")
 parser.add\_argument("--device", default="cuda", help="training device")
 parser.add\_argument("-b", "--batch-size", default=32, type=int)
 parser.add\_argument("--epochs", default=30, type=int, metavar="N",
 help="number of total epochs to train")

 parser.add\_argument('--lr', default=0.0001, type=float, help='initial learning rate')
 parser.add\_argument('--momentum', default=0.9, type=float, metavar='M',
 help='momentum')
 parser.add\_argument('--wd', '--weight-decay', default=1e-4, type=float,
 metavar='W', help='weight decay (default: 1e-4)',
 dest='weight\_decay')
 parser.add\_argument('--print-freq', default=10, type=int, help='print frequency')
 parser.add\_argument('--resume', default='', help='resume from checkpoint')
 parser.add\_argument('--start-epoch', default=0, type=int, metavar='N',
 help='start epoch')
 # 是否使用混合精度训练
 parser.add\_argument("--amp", default=False, type=bool,
 help="Use torch.cuda.amp for mixed precision training")

 args = parser.parse\_args()

 return args

data-path:定义数据集的根目录(VOCdevkit)所在的文件夹

num-classes:检测目标类别数(不包含背景)。

aux:是否使用aux_classifier。

device:使用cpu还是gpu训练,默认是cuda。

batch-size:BatchSize设置。

epochs:epoch的个数。

lr:学习率。

resume:继续训练时候,选择用的模型。

start-epoch:起始的epoch,针对再次训练时,可以不需要从0开始。

amp:是否使用torch的自动混合精度训练。

数据增强

增强调用transforms.py中的方法。

训练集的增强如下:

class SegmentationPresetTrain:
 def \_\_init\_\_(self, base\_size, crop\_size, hflip\_prob=0.5, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
 # 随机Resize的最小尺寸
 min\_size = int(0.5 * base\_size)
 # 随机Resize的最大尺寸
 max\_size = int(2.0 * base\_size)
 # 随机Resize增强。
 trans = [T.RandomResize(min\_size, max\_size)]
 if hflip\_prob > 0:
 #随机水平翻转
 trans.append(T.RandomHorizontalFlip(hflip\_prob))
 trans.extend([
 #随机裁剪
 T.RandomCrop(crop\_size),
 T.ToTensor(),
 T.Normalize(mean=mean, std=std),
 ])
 self.transforms = T.Compose(trans)

 def \_\_call\_\_(self, img, target):
 return self.transforms(img, target)

训练集增强,包括随机Resize、随机水平翻转、随即裁剪。

验证集增强:

class SegmentationPresetEval:
 def \_\_init\_\_(self, base\_size, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
 self.transforms = T.Compose([
 T.RandomResize(base\_size, base\_size),
 T.ToTensor(),
 T.Normalize(mean=mean, std=std),
 ])

 def \_\_call\_\_(self, img, target):
 return self.transforms(img, target)

验证集的增强比较简单,只有随机Resize。

Main方法

对Main方法,我做了一些修改,修改的代码如下:

 #定义模型,并加载预训练
 model = fcn\_resnet50(pretrained=True)
 # 默认classes是21,如果不是21,则要修改类别。
 if num\_classes != 21:
 model.classifier[4] = torch.nn.Conv2d(512, num\_classes, kernel\_size=(1, 1), stride=(1, 1))
 model.aux\_classifier[4] = torch.nn.Conv2d(256, num\_classes, kernel\_size=(1, 1), stride=(1, 1))
 print(model)
 model.to(device)
 # 如果有多张显卡,则使用多张显卡
 if torch.cuda.device\_count() > 1:
 print("Let's use", torch.cuda.device\_count(), "GPUs!")
 model = torch.nn.DataParallel(model)

模型,我改为pytorch官方的模型了,如果能使用官方的模型尽量使用官方的模型。

默认类别是21,如果不是21,则要修改类别。

检测系统中是否有多张卡,如果有多张卡则使用多张卡不能浪费资源。

如果不想使用所有的卡,而是指定其中的几张卡,可以使用:

os.environ['CUDA\_VISIBLE\_DEVICES'] = '0,1'

也可以在DataParallel方法中设定:

model = torch.nn.DataParallel(model,device\_ids=[0,1])

如果使用了多显卡,再使用模型的参数就需要改为model.module.xxx,例如:

  params = [p for p in model.module.aux\_classifier.parameters() if p.requires\_grad]
 params\_to\_optimize.append({"params": params, "lr": args.lr * 10})

上面的都完成了就可以开始训练了,如下图:

938e5d2cc4daeed7805ea8f3d18f2958 - 全卷积网络(FCN)实战:使用FCN实现语义分割

测试

在开始测试之前,我们还要获取到调色板,新建脚本get_palette.py,代码如下:

import json
import numpy as np
from PIL import Image
# 读取mask标签
target = Image.open("./2007\_001288.png")
# 获取调色板
palette = target.getpalette()

palette = np.reshape(palette, (-1, 3)).tolist()
print(palette)
# 转换成字典子形式
pd = dict((i, color) for i, color in enumerate(palette))

json\_str = json.dumps(pd)
with open("palette.json", "w") as f:
 f.write(json\_str)

选取一张mask,然后使用getpalette方法获取,然后将其转为字典的格式保存。

接下来,开始预测部分,新建predict.py,插入以下代码:

import os
import time
import json
import torch
from torchvision import transforms
import numpy as np
from PIL import Image
from torchvision.models.segmentation import fcn\_resnet50

导入程序需要的包文件,然在mian方法中:

def main():
 aux = False # inference time not need aux\_classifier
 classes = 20
 weights\_path = "./save\_weights/model\_5.pth"
 img\_path = "./2007\_000123.jpg"
 palette\_path = "./palette.json"
 assert os.path.exists(weights\_path), f"weights {weights\_path} not found."
 assert os.path.exists(img\_path), f"image {img\_path} not found."
 assert os.path.exists(palette\_path), f"palette {palette\_path} not found."
 with open(palette\_path, "rb") as f:
 pallette\_dict = json.load(f)
 pallette = []
 for v in pallette\_dict.values():
 pallette += v
  • 定义是否需要aux_classifier,预测不需要aux_classifier,所以设置为False。
  • 设置类别为20,不包括背景。
  • 定义权重的路径。
  • 定义调色板的路径。
  • 读去调色板。

接下来,是加载模型,单显卡训练出来的模型和多显卡训练出来的模型加载有区别,我们先看单显卡训练出来的模型如何加载。

   model = fcn\_resnet50(num\_classes=classes+1)
 print(model)
 # 单显卡训练出来的模型,加载
 # delete weights about aux\_classifier
 weights\_dict = torch.load(weights\_path, map\_location='cpu')['model']
 for k in list(weights\_dict.keys()):
 if "aux\_classifier" in k:
 del weights\_dict[k]

 # load weights
 model.load\_state\_dict(weights\_dict)
 model.to(device)

定义模型fcn_resnet50,num_classes设置为类别+1(背景)

加载训练好的模型,并将aux_classifier删除。

然后加载权重。

再看多显卡的模型如何加载

 # create model
 model = fcn\_resnet50(num\_classes=classes+1)
 model = torch.nn.DataParallel(model)
 # delete weights about aux\_classifier
 weights\_dict = torch.load(weights\_path, map\_location='cpu')['model']
 print(weights\_dict)
 for k in list(weights\_dict.keys()):
 if "aux\_classifier" in k:
 del weights\_dict[k]
 # load weights
 model.load\_state\_dict(weights\_dict)
 model=model.module
 model.to(device)

定义模型fcn_resnet50,num_classes设置为类别+1(背景),将模型放入DataParallel类中。

加载训练好的模型,并将aux_classifier删除。

加载权重。

执行torch.nn.DataParallel(model)时,model被放在了model.module,所以model.module才真正需要的模型。所以我们在这里将model.module赋值给model。

接下来是图像数据的处理

 # load image
 original\_img = Image.open(img\_path)

 # from pil image to tensor and normalize
 data\_transform = transforms.Compose([transforms.Resize(520),
 transforms.ToTensor(),
 transforms.Normalize(mean=(0.485, 0.456, 0.406),
 std=(0.229, 0.224, 0.225))])
 img = data\_transform(original\_img)
 # expand batch dimension
 img = torch.unsqueeze(img, dim=0)

加载图像。

对图像做Resize、标准化、归一化处理。

使用torch.unsqueeze增加一个维度。

完成图像的处理后,就可以开始预测了。

 model.eval() # 进入验证模式
 with torch.no\_grad():
 # init model
 img\_height, img\_width = img.shape[-2:]
 init\_img = torch.zeros((1, 3, img\_height, img\_width), device=device)
 model(init\_img)

 t\_start = time\_synchronized()
 output = model(img.to(device))
 t\_end = time\_synchronized()
 print("inference+NMS time: {}".format(t\_end - t\_start))

 prediction = output['out'].argmax(1).squeeze(0)
 prediction = prediction.to("cpu").numpy().astype(np.uint8)
 np.set\_printoptions(threshold=sys.maxsize)
 print(prediction.shape)
 mask = Image.fromarray(prediction)
 mask.putpalette(pallette)
 mask.save("test\_result.png")

将预测后的结果保存到test_result.png中。查看运行结果:

原图:

5a5c1339fec2c2492c01cd8597fef889 - 全卷积网络(FCN)实战:使用FCN实现语义分割

结果:

e6f0223a4b5814247004fe756319ebc7 - 全卷积网络(FCN)实战:使用FCN实现语义分割

打印出来的数据:

f3ba29e2fdeacdad581e09ca2e2582fc - 全卷积网络(FCN)实战:使用FCN实现语义分割

类别列表:

{
 "aeroplane": 1,
 "bicycle": 2,
 "bird": 3,
 "boat": 4,
 "bottle": 5,
 "bus": 6,
 "car": 7,
 "cat": 8,
 "chair": 9,
 "cow": 10,
 "diningtable": 11,
 "dog": 12,
 "horse": 13,
 "motorbike": 14,
 "person": 15,
 "pottedplant": 16,
 "sheep": 17,
 "sofa": 18,
 "train": 19,
 "tvmonitor": 20
}

从结果来看,已经预测出来图像上的类别是“train”。

总结

这篇文章的核心内容是讲解如何使用FCN实现图像的语义分割。

在文章的开始,我们讲了一些FCN的结构和优缺点。然后,讲解了如何读取数据集。接下来,告诉大家如何实现训练。最后,是测试以及结果展示。希望本文能给大家带来帮助。

完整代码:https://download.csdn.net/download/hhhhhhhhhhwwwwwwwwww/83778007

点击关注,第一时间了解华为云新鲜技术~

转载请注明:xuhss » 全卷积网络(FCN)实战:使用FCN实现语义分割

喜欢 (6)

您必须 登录 才能发表评论!