使用Matplotlib库可视化COCO数据集格式常用函数 (二)-蒲公英云

使用Matplotlib库可视化COCO数据集格式常用函数 (二)

文章目录

前言
可视化标注文件中的检测框
可视化标注文件中的掩码
可视化预测文件的检测框和掩码
同时可视化原始图像，标注结果，预测结果

前言

在文章(一)中我们主要介绍了plt包绘制图像的原理和基本使用方法，下面我们借用pycocotools包介绍COCO格式数据的可视化。

pycocotools包是官方给出的包，提供了很多针对COCO格式的API，对于pycocotools包的详细实现可以查看coco官方api详解。

可视化标注文件中的检测框

为了可视化标注文件中的检测框，需要执行如下步骤：

首先根据标注的json文件构建一个COCO对象；
调用coco.getAnnoIds(image_id)函数得到某张图片所有标注的id；
调用coco.loadAnno(anno_id)函数加载所有的标注；

使用cv2或者plt可视化检测框；
代码如下,代码链接：

import os
import cv2

from pycocotools.coco import COCO

json_file = r’F:/dataset/coco/2017/annotations/instances_val2017.json’
dataset_dir = r’F:/dataset/coco/2017/val2017/‘
coco = COCO(json_file)
catIds = coco.getCatIds(catNms=[‘person’]) # catIds=1 表示人这一类
imgIds = coco.getImgIds(catIds=catIds ) # 图片id，许多值
for i in range(len(imgIds)):

img = coco.loadImgs(imgIds[i])[0]
image = cv2.imread(dataset_dir + img['file_name'])
annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
annos = coco.loadAnns(annIds)
bbox = annos[0]['bbox']
x, y, w, h = bbox
anno_image = cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 255), 2) 
# 参数为(显示的图片名称，要显示的图片)  必须加上图片名称，不然会报错
cv2.imshow('demo', anno_image)
cv2.waitKey(5000)

可视化标注文件中的掩码

可视化标注文件中的掩码比较容易，步骤和可视化检测框类似，只不过显示掩码时可以直接调用coco.showAnns()实现。

代码如下：

import os
from pycocotools.coco import COCO
from skimage import io
from matplotlib import pyplot as plt
json_file = '../datasets/coco/annotations/instances_val2017.json'
dataset_dir = '../datasets/coco/images/val2017/'
coco = COCO(json_file)
catIds = coco.getCatIds(catNms=['person']) # catIds=1 表示人这一类
imgIds = coco.getImgIds(catIds=catIds ) # 图片id，许多值
for i in range(len(imgIds)):
    img = coco.loadImgs(imgIds[i])[0]
    I = io.imread(dataset_dir + img['file_name'])
    plt.axis('off')
    plt.imshow(I) #绘制图像，显示交给plt.show()处理
    annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
    anns = coco.loadAnns(annIds)
    coco.showAnns(anns)
    plt.show() #显示图像
    plt.savefig(mask_dir + r'/' + img['file_name'])
    plt.clf()  # 清除画布

可视化预测文件的检测框和掩码

一般detectron2预测的结果会保存在一个result.json文件中，为了可视化预测文件的检测框和掩码，我们可以调用coco.loadRes()函数将预测的json文件生成coco对象，然后和上面介绍的方法一样，调用官方API进行可视化。

代码如下：

# 根据真实标签json文件生成coco对象
coco = COCO(annotation_file)
# 根据预测json文件生成coco对象
coco_res = coco.loadRes('C:/Users/hejianfei/Desktop/inference/results12-162.json')

同时可视化原始图像，标注结果，预测结果

下面给出一个同时可视化原始图像，标注结果，预测结果的例子。

# 测试集标注json文件
annotation_file = 'F:/9.4Data/ski10/annotations/ski10_val.json'
# 测试集图片文件夹
image_folder = 'F:/9.4Data/ski10/images-test'
# 保存结果文件夹
save_folder = 'F:/9.4Data/ski10/visualize-res-val'
# 根据真实标签json文件生成coco对象
coco = COCO(annotation_file)
# 根据预测json文件生成coco对象
coco_res = coco.loadRes('C:/Users/hejianfei/Desktop/inference/results12-162.json')
# 如果文件夹不存在，则创建
if not os.path.exists(save_folder):
    os.makedirs(save_folder)
# 获取所有的测试图像id
image_ids = coco.getImgIds()
print('本次处理的图片数为' + str(len(image_ids)))
# 依次处理每一张图片：读取显示图片，显示每张图片对应的标签
for image_id in image_ids:
    image = coco.loadImgs(image_id)[0]
    I = io.imread(image_folder + r'/' + image['file_name'])
    plt.axis('off')
    # 加载并显示原图
    plt.subplot(1, 3, 1)
    plt.title('raw-image')
    plt.imshow(I)  # 绘制图像，显示交给plt.show()处理
    # 加载并显示正确的anno
    plt.subplot(1, 3, 2)
    plt.title('true-label')
    plt.imshow(I)
    annIds = coco.getAnnIds(imgIds=image_id, iscrowd=None)
    anns = coco.loadAnns(annIds)
    coco.showAnns(anns, draw_bbox=False)
    # 加载并显示预测的anno
    plt.subplot(1, 3, 3)
    plt.title('pred-label')
    plt.imshow(I)
    annIds = coco_res.getAnnIds(imgIds=image_id, iscrowd=None)
    anns = coco_res.loadAnns(annIds)
    coco_res.showAnns(anns, draw_bbox=False)
    # 保存结果
    plt.savefig(save_folder + '/{}.jpg'.format(str(image_id)))
    # plt.show()
    plt.clf()  # 清除画布