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

浅浅的花香味﹌ 2024-05-08 06:39 5阅读 0赞

#### 文章目录 ####

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

## 前言 ##

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

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

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

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

*  首先根据标注的json文件构建一个COCO对象；
 *  调用`coco.getAnnoIds(image_id)`函数得到某张图片所有标注的id；
 *  调用`coco.loadAnno(anno_id)`函数加载所有的标注；
 *  使用cv2或者plt可视化检测框；  
    代码如下,[代码链接][Link 1]：

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()  # 清除画布

结果如下  
![在这里插入图片描述][watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_18_color_FFFFFF_t_70_g_se_x_16]

[coco_api]: https://blog.csdn.net/qq_43705697/article/details/122339511
[Link 1]: https://blog.csdn.net/qq_29631521/article/details/103166116
[watermark_type_d3F5LXplbmhlaQ_shadow_50_text_Q1NETiBA6YGH5Yiw5Z2O5bCx5b6X6L-I6L-H5Y67_size_18_color_FFFFFF_t_70_g_se_x_16]: https://image.dandelioncloud.cn/pgy_files/images/2024/05/08/bd661b1c7ea843f0a41e2e5e73cd2cb0.png