欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

PASCAL VOC图像语义分割(Pytorch)

程序员文章站 2022-07-05 11:05:39
...

PASCAL VOC 2012数据集介绍

Pascal VOC的三个主要物体识别竞赛是:分类,检测和分割(classification, detection, and segmentation)。对于分割任务, VOC2012的trainval包含2007-2011年的所有对应图片, test只包含2008-2011。trainval有 2913张图片共6929个物体。
数据集将会放置在D:\deeplearning\VOCdevkit\VOC2012路径下。进入D:\deeplearning\VOCdevkit\VOC2012路径后,我们可以获取数据集的不同组成部分。其中ImageSets\Segmentation路径包含了指定训练和测试样本的文本文件,而JPEGImages和SegmentationClass路径下分别包含了样本的输入图像和标签。这里的标签也是图像格式,其尺寸和它所标注的输入图像的尺寸相同。标签中颜色相同的像素属于同一个语义类别。

Pascal VOC2012数据集主要是针对视觉任务中监督学习提供标签数据,它有二十个类别(背景为第21类)
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
主要有四个大类别,分别是人、常见动物、交通车辆、室内家具用品。主要为图像分类、对象检测识别、图像分割三类任务服务。

实验过程

(1)读取图像数据, 将原图和分割标准图像路径改为数据存放路径。
(2)将输入图像转换为模型需要的batchg。原始输入图像大小为Ci×Hi×Wi,其中Ci为图像通道数量(RGB图像Ci=3),Hi为图像的高,Wi为图像的宽。而预训练模型期望的输入是RGB图像的mini-batch:(batch_size, 3, H, W)。所以我们要将输入图像变换为Ni×Ci×Hi×Wi,其中Ni为batch大小。此外,Hi和Wi不能低于224,图像的像素值必须在范围[0,1]间,所以用均值mean=[0.485, 0.456, 0.406]和方差std=[0.229, 0.224, 0.225]进行归一化处理
(3)加载预训练的图像分割全卷积网络模型(FCN-ResNet101),并对输入图像数据进行预测。加载预训练模型主要是通过models.segmentation.fcn_resnet101(pretrained=True)进行。Torch的FCN-ResNet101语义分割模型是在COCO 2017训练集上的一个子集训练得到的,相当于PASCAL VOC数据集,支持20个类别(背景为第21个类别)。
(4)将输出标签转换为VOC分类标签格式。将得到的灰度图变换到调色板图像模式,预测图像的调色板信息从标准分割图像中提取。获取ground truth标签图像调色板以用于预测图像。
(5)保存分割后的图像,即模型输出的分割图像。
(6)输出对比图像,依次为:原图、分割标准图像、模型分割后的图像、融合图像。融合图像为原图与模型输出的分割图像的融合。

具体代码
from PIL import Image
import numpy as np 
import torchvision.utils as vutils
import torchvision.models as models
import torchvision.transforms as T 
import matplotlib.pyplot as plt
# 1.
# load image data
# 读图像数据, 将原图和分割标准图像路径改为数据存放路径
img_path = 'D:\\hxq\\deeplearning\\VOCdevkit\\VOC2012\\JPEGImages\\2007_000033.jpg'
label_path = 'D:\\hxq\\deeplearning\\VOCdevkit\\VOC2012\\SegmentationClass\\2007_000033.png'
img = Image.open(img_path)
label_img = Image.open(label_path)
# 2
# get input data batch 将输入图像变换为模型需要的batchg
#预训练模型期望的输入是RGB图像的mini-batch:(batch_size, 3, H, W),
#并且H和W不能低于224。图像的像素值必须在范围[0,1]间,
#并且用均值mean=[0.485, 0.456, 0.406]和方差std=[0.229, 0.224, 0.225]进行归一化。
test_transform = T.Compose(
    [T.ToTensor(),
    T.Normalize(mean = [0.485,0.456,0.406],std=[0.229,0.224,0.225])]
)
input_tensor = test_transform(img)
input_batch = input_tensor.unsqueeze(0)
# [Ci x Hi x Wi]->[Ni x Ci x Hi x Wi] 彩色图像3通道,变换到batch 
#Ni -> the batch size 
#Ci -> the number of channels (which is 3) 图像通道
#Hi -> the height of the image 高
#Wi -> the width of the image 宽
# 3
# load pretrained segmentation model and predict input data batch
# wait model downloading for the firt time.
#预训练模型可以通过设置pretrained=True来构建:
seg_model = models.segmentation.fcn_resnet101(pretrained=True)
# 上面使用了全卷积网络 FCN, 如果要使用DeepLabV3, 把上面注释掉,下面取消注释
#seg_model = models.segmentation.deeplabv3_resnet101(pretrained=True)
seg_model.eval()
output = seg_model(input_batch)['out'][0]
output_prediction = output.argmax(0)
# 4
# Transform output label to VOC segmentation label format
label_prediction = output_prediction.numpy().astype(np.uint8)
img_prediction = T.ToPILImage()(label_prediction).convert('P')
# change gray image to image with color palette 
# 得到的灰度图变换到调色板图像模式, 预测图像的调色板信息从标准分割图像中提取
color_palette=label_img.getpalette()
img_prediction.putpalette(color_palette)
# get ground truth label image color palette for prediction image
# 5
# Only save output file /仅保存分割后的图像,不需要输出对比图像可视化到此为止
img_prediction.save('output.png')
# 6
# visualization /可视化,输出对比图像依次为:原图、分割标准图像、模型分割后的图像、融合图像
mini_batch = []
# orginal image 原图
img_tensor=T.functional.to_tensor(img)
mini_batch.append(img_tensor)
# ground truth label image 分割标准图像
label_img_tensor = T.functional.to_tensor(label_img.convert('RGB'))
mini_batch.append(label_img_tensor)
# prediction label image 模型分割后的图像
img_prediction_rgb = img_prediction.convert('RGB')
img_prediction_tensor = T.functional.to_tensor(img_prediction_rgb)
mini_batch.append(img_prediction_tensor)
# blending image 融合图像
blend_img = Image.blend(img,img_prediction_rgb,alpha=0.5)
blend_img_tensor = T.functional.to_tensor(blend_img)
mini_batch.append(blend_img_tensor)
# Show images with matplotlib /用matplotlib显示图像
grid_img = vutils.make_grid(mini_batch,padding=3,pad_value=1)
plt.axis('off')
plt.imshow(grid_img.permute(1,2,0))
plt.show()
# Save result image to file /保存对比图像到文件
vutils.save_image(mini_batch,'result.png')
实验结果

实验结果对比图,对比图像依次为:原图、分割标准图像、模型分割后的图像、融合图像(模型分割后的图像与原图的融合)
PASCAL VOC图像语义分割(Pytorch)
PASCAL VOC图像语义分割(Pytorch)
PASCAL VOC图像语义分割(Pytorch)
通过与分割标准图像的对比,可以发现该模型的输出分割图像与分割标准图像几乎一致,同时模型的输出分割图像与原图也较好的融合,说明该模型具有较好的准确性。此外,从输入图像大小来看,该模型可以输入任意大小的图像,并输出相同大小的已经标签好的分割图像。由于是针对PASCAL VOC数据集图像进行的分割,PASCAL VOC数据集中只支持20个类别(背景为第21个类别),所以在分割时,遇到不在20个类别中的事物都将其标为背景。但总体来说,该模型对PASCAL VOC数据集的图像分割达到了较高准确率。

代码参考:
https://zhuanlan.zhihu.com/p/99250652