[AI达人特训营第三期]寻找被放走的鸽子
本项目使用OCRNet实现了鸽子语义分割模型,能够达到91.56%的mIoU,是一个比较成功的语义分割案例。另外本项目还提供了一个Paddle-Inference部署案例供大家学习参考。
★★★ 本文源自AlStudio社区精品项目,【点击此处】查看更多精品内容 >>>
本项目使用OCRNet实现了鸽子语义分割模型,各项指标如下:
模型名称 | mIoU | Acc | Kappa | Dice |
---|---|---|---|---|
OCRNet | 0.9156 | 0.9929 | 0.9085 | 0.9542 |
整体的实现效果是不错的,另外本项目提供了一个Paddle-Inference-Demo供大家参考。
一、项目背景
对于语义分割任务来说,其两大关键是:分辨率和上下文。
- 语义分割是一个密集像素预测任务,因此空间分辨率很重要。
- 而像素本身不具备语义,它的语义由其图像整体或目标区域决定,因此它对上下文高度依赖,一个像素位置的上下文指的是它周围的像素位置。
OCRNet论文的主要思想也就是像素的类别标签是由它所在的目标的类别标签决定的。主要思路是利用目标区域表示来增强其像素的表示。与之前的考虑上下文关系的方法不同的是,之前的方法考虑的是上下文像素之间的关系,没有显示利用目标区域的特征。
其中,粉红色虚线框内为形成的软对象区域(Soft Object Regions),紫色虚线框中为物体区域表示(Object Region Representations),橙色虚线框中为对象上下文表示和增强表示。
第一步: 将上下文像素划分为一组软对象区域,每个soft object regions对应一个类,即从深度网络(backbone)计算得到的粗软分割(粗略的语义分割结果)。这种划分是在ground-truth分割的监督下学习的。根据网络中间层的特征表示估测粗略的语义分割结果作为 OCR 方法的一个输入,即结构图中粉红色框内的Soft Object Regions。
第二步: 根据粗略的语义分割结果(soft object regions)和网络最深层输出的像素特征(Pixel Representations)表示计算出 K 组向量,即物体区域表示(Object Region Representations),其中每一个向量对应一个语义类别的特征表示。
第三步: 这一步可以再细分为两个步骤。
- 计算网络最深层输出的像素特征表示(Pixel Representations)与计算得到的物体区域特征表示(Object Region Representation)之间的关系矩阵,然后根据每个像素和物体区域特征表示在关系矩阵中的数值把物体区域特征加权求和,得到最后的物体上下文特征表示 (Object Contextual Representation),即OCR 。
- 当把物体上下文特征表示 OCR 与网络最深层输入的特征表示拼接之后作为上下文信息增强的特征表示(Augmented Representation),可以基于增强后的特征表示预测每个像素的语义类别。
二、数据集简介
该数据集由124张鸽子图片及其标注图组成,现已在AI Studio上开源。链接
三、数据预处理
Step01: 解压数据集
!unzip /home/aistudio/data/data75217/doves.zip -d /home/aistudio/work/doves/
Step02: 将原图和标注图分隔开。
在/home/aistudio/work/data目录下新建images和labels文件夹。
- images:原始图像。
- labels:标注文件。
import os
import random
import cv2
def get_images(image_dir):
"""
取出数据集中的原图和标签图,并重新写入新的文件夹
"""
images = []
labels = []
for image_name in os.listdir(image_dir):
if os.path.isdir(os.path.join(image_dir, image_name)):
img = cv2.imread(os.path.join(os.path.join(image_dir, image_name), 'img.png'))
cv2.imwrite("/home/aistudio/work/data/images/" + image_name + '.png', img)
print("/home/aistudio/work/data/images/" + image_name + '.png')
label = cv2.imread(os.path.join(os.path.join(image_dir, image_name), 'label.png'))
cv2.imwrite("/home/aistudio/work/data/labels/" + image_name + '.png', label)
print("/home/aistudio/work/data/labels/" + image_name + '.png')
if __name__ == "__main__":
get_images("/home/aistudio/work/doves")
四、代码实现
4.1 环境配置
从Github下载PaddleSeg代码。
#!git clone https://github.com/PaddlePaddle/PaddleSeg
执行如下命令,从源码编译安装PaddleSeg包。大家对于PaddleSeg/paddleseg目录下的修改,都会立即生效,无需重新安装。
%cd /home/aistudio/PaddleSeg/
!pip install -r requirements.txt --user
!pip install -v -e .
4.2 数据准备
Step01: 将数据集移动到/home/aistudio/PaddleSeg/data目录下。
!mv /home/aistudio/work/data /home/aistudio/PaddleSeg/data/
Step02: 切分数据。
!python tools/data/split_dataset_list.py /home/aistudio/PaddleSeg/data/data images labels --split 0.85 0.15 0 --format png png
4.3 模型训练
Step01: 图像尺寸分析
通过图像分析,我们可以看到数据集中只有1种不同的尺寸,是[3024, 4032]。
def Image_size(indir):
width_heights = []
for img in os.listdir(indir):
img = cv2.imread(indir + img)
width = img.shape[0]
height = img.shape[1]
if [width, height] not in width_heights: width_heights.append([width, height])
print("数据集中,有{}种不同的尺寸,分别是:".format(len(width_heights)))
for item in width_heights:
print(item)
indir='/home/aistudio/PaddleSeg/data/data/images/' # xml文件所在的目录
Image_size(indir)
Step02: 单卡训练
ERROR1: 标注图是彩色标注图,而不是灰度标注图。
ValueError: (InvalidArgument) If Attr(soft_label) == false, the axis dimension of Input(Label) should be 1.
[Hint: Expected labels_dims[axis] == 1UL, but received labels_dims[axis]:3 != 1UL:1.] (at /paddle/paddle/phi/infermeta/binary.cc:937)
SOLUTION1: 通过以下代码转换成灰度标注图。
import os
import numpy as np
from PIL import Image
from tqdm import tqdm
Origin_SegmentationClass_path = "/home/aistudio/PaddleSeg/data/data/labels"
Out_SegmentationClass_path = "/home/aistudio/PaddleSeg/data/data/mask"
# 对应关系
Origin_Point_Value = np.array([[0, 0, 0], [128, 0, 0]])
Out_Point_Value = np.array([0, 1])
if __name__ == "__main__":
if not os.path.exists(Out_SegmentationClass_path):
os.makedirs(Out_SegmentationClass_path)
#
png_names = os.listdir(Origin_SegmentationClass_path) # 获得图片的文件名
print("正在遍历全部标签。")
for png_name in tqdm(png_names):
if png_name == ".ipynb_checkpoints":
continue
png = Image.open(os.path.join(Origin_SegmentationClass_path, png_name)) # RGB
w, h = png.size
png = np.array(png, np.uint8) # h, w, c
out_png = np.zeros([h, w]) # 灰度 h, w
for map_idx, rgb in enumerate(Origin_Point_Value):
idx = np.where(
(png[..., 0] == rgb[0]) & (png[..., 1] == rgb[1]) & (png[..., 2] == rgb[2]))
out_png[idx] = map_idx
# print("out_png:", out_png.shape)
out_png = Image.fromarray(np.array(out_png, np.uint8)) # 再次转化为Image进行保存
out_png.save(os.path.join(Out_SegmentationClass_path, png_name))
# 统计输出,各个像素点的值的个数
print("正在统计输出的图片每个像素点的数量。")
classes_nums = np.zeros([256], np.int32)
for png_name in tqdm(png_names):
if png_name == ".ipynb_checkpoints":
continue
png_file_name = os.path.join(Out_SegmentationClass_path, png_name)
if not os.path.exists(png_file_name):
raise ValueError("未检测到标签图片%s,请查看具体路径下文件是否存在以及后缀是否为png。" % (png_file_name))
png = np.array(Image.open(png_file_name), np.uint8)
classes_nums += np.bincount(np.reshape(png, [-1]), minlength=256)
print("打印像素点的值与数量。")
print('-' * 37)
print("| %15s | %15s |" % ("Key", "Value"))
print('-' * 37)
for i in range(256):
if classes_nums[i] > 0:
print("| %15s | %15s |" % (str(i), str(classes_nums[i])))
print('-' * 37)
!python tools/data/split_dataset_list.py /home/aistudio/PaddleSeg/data/data images mask --split 0.85 0.15 0 --format png png
!python tools/train.py \
--config configs/ocrnet/ocrnet_hrnetw48_voc12aug_512x512_40k.yml \
--do_eval \
--use_vdl \
--save_interval 100 \
--save_dir output/ocrnet
损失函数如图所示:
4.4 模型评估
训练完成后,大家可以使用评估脚本tools/val.py来评估模型的精度,即对配置文件中的验证数据集进行测试。
!python tools/val.py \
--config configs/ocrnet/ocrnet_hrnetw48_voc12aug_512x512_40k.yml \
--model_path output/ocrnet/best_model/model.pdparams
- [EVAL] #Images: 18 mIoU: 0.9156 Acc: 0.9929 Kappa: 0.9085 Dice: 0.9542
- [EVAL] Class IoU: [0.9926 0.8385]
- [EVAL] Class Precision: [0.9956 0.9281]
- [EVAL] Class Recall: [0.997 0.8968]
4.5 模型预测
我们可以通过tools/predict.py脚本是来进行可视化预测,命令格式如下所示。
!python tools/predict.py \
--config configs/ocrnet/ocrnet_hrnetw48_voc12aug_512x512_40k.yml \
--model_path output/ocrnet/best_model/model.pdparams \
--image_path data/data/images \
--save_dir output/result
整体的检测效果是非常不错的,部分可视化结果如下:
4.6 模型导出
执行如下命令,导出预测模型,保存在output/inference_model目录。
!python tools/export.py \
--config configs/ocrnet/ocrnet_hrnetw48_voc12aug_512x512_40k.yml \
--model_path output/ocrnet/best_model/model.pdparams \
--save_dir output/inference_model
4.7 Paddle-Inference-demo
- 引用 paddle inference 推理库
import paddle.inference as paddle_infer
import cv2
import numpy as np
- 创建配置对象
# 创建 config,并设置推理模型路径
config = paddle_infer.Config("/home/aistudio/PaddleSeg/output/inference_model/model.pdmodel", "/home/aistudio/PaddleSeg/output/inference_model/model.pdiparams")
- 根据Config创建推理对象
predictor = paddle_infer.create_predictor(config)
- 设置模型输入 Tensor
def normalize(img, mean, std):
img = img / 255.0
mean = np.array(mean)[np.newaxis, np.newaxis, :]
std = np.array(std)[np.newaxis, np.newaxis, :]
img -= mean
img /= std
return img
def preprocess(img, img_size):
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]
img = cv2.resize(img, (img_size[0], img_size[1]))
img = img[:, :, ::-1].astype('float32') # bgr -> rgb
img = normalize(img, mean, std)
img = img.transpose((2, 0, 1)) # hwc -> chw
return img[np.newaxis, :]
# 设置输入
img = cv2.imread("/home/aistudio/PaddleSeg/data/data/images/IMG_4676.png")
print(img.shape)
im_shape = img.shape
im_size = [512, 512]
data = preprocess(img, im_size)
# 获取输入的名称
input_names = predictor.get_input_names()
input_tensor = predictor.get_input_handle(input_names[0])
input_tensor.reshape(img.shape)
input_tensor.copy_from_cpu(data.copy())
- 执行推理
predictor.run()
- 获得推理结果
COLORMAP = np.array([[28, 28, 28], [238, 44, 44]])
def segmentation_map_to_image(result, colormap, remove_holes=False):
"""
可视化输出特征图
:param result: 分割模型的输出特征图
:param colormap: 类别到颜色的映射关系 class i -> COLORMAP[i]
:param remove_holes: True->去除输出特征图中的空洞
:return: 可视化后的RGB图像
"""
if len(result.shape) != 2 and result.shape[0] != 1:
raise ValueError(
f"Expected result with shape (H,W) or (1,H,W), got result with shape {result.shape}"
)
if len(np.unique(result)) > colormap.shape[0]:
raise ValueError(
f"Expected max {colormap[0]} classes in result, got {len(np.unique(result))} "
"different output values. Please make sure to convert the network output to "
"pixel values before calling this function."
)
elif result.shape[0] == 1:
result = result.squeeze(0)
result = result.astype(np.uint8)
contour_mode = cv2.RETR_EXTERNAL if remove_holes else cv2.RETR_TREE
mask = np.zeros((result.shape[0], result.shape[1], 3), dtype=np.uint8)
for label_index, color in enumerate(colormap):
label_index_map = result == label_index
label_index_map = label_index_map.astype(np.uint8) * 255
contours, hierarchies = cv2.findContours(
label_index_map, contour_mode, cv2.CHAIN_APPROX_SIMPLE
)
cv2.drawContours(
mask,
contours,
contourIdx=-1,
color=color.tolist(),
thickness=cv2.FILLED,
)
return mask
output_names = predictor.get_output_names()
output_tensor = predictor.get_output_handle(output_names[0])
output_data = output_tensor.copy_to_cpu()
output_mask = segmentation_map_to_image(output_data, COLORMAP)
output_mask = cv2.resize(output_mask, (im_shape[1], im_shape[0]))
cv2.imwrite("/home/aistudio/work/output.jpg", output_mask)
print(output_mask.shape)
可视化结果如下:
五、总结提高
本项目任务是使用OCRNet去实现鸽子分割的任务,整体上来说任务是相对较简单的,能够达到91.56%的mIoU,大家可能遇到的问题是训练的过程比较漫长。
另外本项目任务还给大家提供了一个简单的Paddle-Inference-Demo供大家参考。
还可能改进的几个点:
- 数据集较少,大家时间充裕的话可以再搜集一些鸽子的图片,标注好后加入进来一起训练。
- 使用轻量化的网络看是否能够在不损失太多精度的情况下,实现更低的推理延时。
作者简介:Submerge. 江苏某大学大三学生 人工智能专业 主页链接 欢迎互关!
飞桨导师:刘建建 JavaRoom 在此感谢。
此文章为转载
原文链接
更多推荐
所有评论(0)