COCO数据集bbox和instance_masks可视化

最新推荐文章于 2023-11-24 16:28:13 发布

原创最新推荐文章于 2023-11-24 16:28:13 发布 · 3.5k 阅读

15 ·

CC 4.0 BY-SA版权

学习笔记专栏收录该内容

2 篇文章

订阅专栏

本文详细介绍如何使用Python和pycocotools库对COCO数据集进行物体检测框和实例mask的可视化，包括随机选择图像进行物体框的绘制及实例mask的显示。

部署运行你感兴趣的模型镜像

参考博客：COCO数据集可视化 + 物体检测标注格式-coco可视化

1、物体框可视化

from pycocotools.coco import COCO
import cv2 as cv
import os
import numpy as np
import random

img_path = 'val2017/' # 把图片直接放在同一文件夹下
annFile = 'instances_val2017.json' # 同样

coco = COCO(annFile)
cats = coco.loadCats(coco.getCatIds())

catIds = coco.getCatIds(catNms=['yuanhuan', 'yellow', 'green', 'red', 'indicator'])
imgIds = coco.getImgIds(catIds=catIds)
img_list = os.listdir(img_path)
for i in range(len(img_list)):
    img = coco.loadImgs(imgIds[np.random.randint(0, len(imgIds))])[0]
    image_name = img['file_name']

    # 加了catIds就是只加载目标类别的anno，不加就是图像中所有的类别anno
    annIds = coco.getAnnIds(imgIds=img['id'], catIds=[], iscrowd=None)
    anns = coco.loadAnns(annIds)
    
    coco.showAnns(anns)
    
    coordinates = []
    img_raw = cv.imread(os.path.join(img_path, image_name))
    for j in range(len(anns)):
        x1 = int(anns[j]['bbox'][0])
        y1 = int(anns[j]['bbox'][1] + anns[j]['bbox'][3])
        x2 = int(anns[j]['bbox'][0] + anns[j]['bbox'][2])
        y2 = int(anns[j]['bbox'][1])
        cv.rectangle(img_raw,
                      (x1, y1),
                      (x2, y2),
                      (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)),
                      2)
    cv.imshow('image', img_raw)
    cv.waitKey()

样例：

2、实例mask可视化

from pycocotools.coco import COCO
from skimage import io
from matplotlib import pyplot as plt
import cv2 as cv

dataset_dir = 'val2017/'
json_file = 'instances_val2017.json'
coco = COCO(json_file)
# 展示coco的所有类别
print(coco.cats)
# catIds = coco.getCatIds(catNms=['cat', 'person']) # catIds=1,2 表示同时含有人和猫这一类
catIds = coco.getCatIds(catNms=['person']) # catIds=1 表示人这一类
# catIds = coco.getCatIds(catNms=['yuanhuan', 'yellow', 'green', 'red', 'indicator']) # 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'], iscrowd=None)
    # 加了catIds就是只加载目标类别的anno，不加就是图像中所有的类别anno
    # annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
    anns = coco.loadAnns(annIds)
    coco.showAnns(anns)
    plt.savefig('ans.png')
    plt.clf()
    ans = cv.imread('ans.png')
    try:
        cv.imshow('mask', ans)
        cv.waitKey()
    except:
        pass