COCO数据集可视化程序(包括bbox和segmentation)

出于项目需要，需要对COCO数据集做可视化，但网上没有比较好用的可视化方法，于是自己写了一个。分为三部分：1.标注框bbox的单独可视化，2.分割标注segmentation的单独可视化，3.bbox和segmentation的联合可视化。
我的数据集格式是标准的COCO数据集格式，可以百度即可。

1.标注框bbox的单独可视化

1.1 利用opencv可视化

一开始打算用COCO API实现bbox的可视化，但是发现API不能直接可视化bbox，于是参照了COCO数据集可视化，自己写了一个bbox的可视化方法。

import json
import os, cv2

train_json = 'C:/Users/wang/Desktop/train/train_33485.json'
train_path = 'C:/Users/wang/Desktop/train/jpeg/'

def visualization_bbox1(num_image, json_path,img_path):# 需要画的第num副图片， 对应的json路径和图片路径
    with open(json_path) as annos:
        annotation_json = json.load(annos)

    print('the annotation_json num_key is:',len(annotation_json))  # 统计json文件的关键字长度
    print('the annotation_json key is:', annotation_json.keys()) # 读出json文件的关键字
    print('the annotation_json num_images is:', len(annotation_json['images'])) # json文件中包含的图片数量

    image_name = annotation_json['images'][num_image - 1]['file_name']  # 读取图片名
    id = annotation_json['images'][num_image - 1]['id']  # 读取图片id

    image_path = os.path.join(img_path, str(image_name).zfill(5)) # 拼接图像路径
    image = cv2.imread(image_path, 1)  # 保持原始格式的方式读取图像
    num_bbox = 0  # 统计一幅图片中bbox的数量

    for i in range(len(annotation_json['annotations'][::])):
        if  annotation_json['annotations'][i-1]['image_id'] == id:
            num_bbox = num_bbox + 1
            x, y, w, h = annotation_json['annotations'][i-1]['bbox']  # 读取边框
            image = cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 255), 2)

    print('The unm_bbox of the display image is:', num_bbox)

    # 显示方式1：用plt.imshow()显示
    # plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) #绘制图像，将CV的BGR换成RGB
    # plt.show() #显示图像

    # 显示方式2：用cv2.imshow()显示
    cv2.namedWindow(image_name, 0)  # 创建窗口
    cv2.resizeWindow(image_name, 1000, 1000) # 创建500*500的窗口
    cv2.imshow(image_name, image)
    cv2.waitKey(0)

if __name__ == "__main__":
   visualization_bbox1(199, train_json, train_path)

后来发现可以用CV2和plot两种方法显示图片。唯一区别就是二者需要的RGB顺序不一样，可以借助 cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 进行转换。

1.2 利用coco API读取json

在后面写分割可视化时，仔细研究了一下coco API，发现其实可以借助coco API来读json数据进行读取，方便快捷不出错。可以参考COCO数据集及COCOAPI中写的“COCO类定义的10个方法”中的前7个。有了这7个方法，可以非常方便的获取json文件中的类别，图片，标注信息。

import cv2
import random
import json, os
from pycocotools.coco import COCO
from skimage import io
from matplotlib import pyplot as plt

train_json = 'C:/Users/wang/Desktop/train/train_33485.json'
train_path = 'C:/Users/wang/Desktop/train/jpeg/'

def visualization_bbox2(num_image, json_path, img_path):
    coco = COCO(json_path)

    list_imgIds = coco.getImgIds(catIds=catIds ) # 获取含有该给定类别的所有图片的id
    img = coco.loadImgs(list_imgIds[num_image-1])[0]  # 获取满足上述要求，并给定显示第num幅image对应的dict
    image = io.imread(img_path + img['file_name'])  # 读取图像
    image_name =  img['file_name'] # 读取图像名字
    image_id = img['id'] # 读取图像id
    for i in range(len(img_annIds)):
        x, y, w, h = img_anns[i-1]['bbox']  # 读取边框
        image = cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 255), 2)

    plt.rcParams['figure.figsize'] = (20.0, 20.0)
    # 此处的20.0是由于我的图片是2000*2000，目前还没去研究怎么利用plt自动分辨率。
    plt.imshow(image)
    plt.show()

2.分割标注segmentation的单独可视化

coco API中有一个方法 showAnns() 可以方便的完成分割标注的显示，因此直接利用coco API做显示。

import cv2
import random
import json, os
from pycocotools.coco import COCO
from skimage import io
from matplotlib import pyplot as plt

train_json = 'C:/Users/wang/Desktop/train/train_33485.json'
train_path = 'C:/Users/wang/Desktop/train/jpeg/'

def visualization_seg(num_image, json_path, img_path, str = ' '):
    # 需要画图的是第num副图片, 对应的json路径和图片路径,
    # str = ' '为类别字符串，输入必须为字符串形式 'str'，若为空，则返回所有类别id
    coco = COCO(json_path)
    
    catIds = coco.getCatIds(catNms=['str']) # 获取指定类别 id
    
    imgIds = coco.getImgIds(catIds=catIds ) # 获取图片i
    img = coco.loadImgs(imgIds[num_image-1])[0]  # 加载图片,loadImgs() 返回的是只有一个内嵌字典元素的list, 使用[0]来访问这个元素
    image = io.imread(train_path + img['file_name'])
   
    annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
    anns = coco.loadAnns(annIds)
   
    # 读取在线图片的方法
    # I = io.imread(img['coco_url'])
    
    plt.imshow(image) 
    coco.showAnns(anns)
    plt.show() 
if __name__ == "__main__":
   visualization_seg(136, train_json, train_path, 'person')

3.bbox和segmentation的联合可视化

最后就是在同一张图片上把bbox和segmentation同时标注出来。

import cv2
import random
import json, os
from pycocotools.coco import COCO
from skimage import io
from matplotlib import pyplot as plt

train_json = 'C:/Users/wang/Desktop/train/train_33485.json'
train_path = 'C:/Users/wang/Desktop/train/jpeg/'

def visualization_bbox_seg(num_image, json_path, img_path, *str):# 需要画图的是第num副图片， 对应的json路径和图片路径

    coco = COCO(json_path)

    if len(str) == 0:
        catIds = []
    else:
        catIds = coco.getCatIds(catNms = [str[0]])  # 获取给定类别对应的id 的dict（单个内嵌字典的类别[{}]）
        catIds = coco.loadCats(catIds)[0]['id'] # 获取给定类别对应的id 的dict中的具体id

    list_imgIds = coco.getImgIds(catIds=catIds ) # 获取含有该给定类别的所有图片的id
    img = coco.loadImgs(list_imgIds[num_image-1])[0]  # 获取满足上述要求，并给定显示第num幅image对应的dict
    image = io.imread(img_path + img['file_name'])  # 读取图像
    image_name =  img['file_name'] # 读取图像名字
    image_id = img['id'] # 读取图像id

    img_annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None) # 读取这张图片的所有seg_id
    img_anns = coco.loadAnns(img_annIds)

    for i in range(len(img_annIds)):
        x, y, w, h = img_anns[i-1]['bbox']  # 读取边框
        image = cv2.rectangle(image, (int(x), int(y)), (int(x + w), int(y + h)), (0, 255, 255), 2)

    plt.rcParams['figure.figsize'] = (20.0, 20.0)
    plt.imshow(image)
    coco.showAnns(img_anns)
    plt.show()

if __name__ == "__main__":
   visualization_bbox_seg(257, train_json, train_path,'snowboard') # 最后一个参数不写就是画出一张图中的所有类别

其他的信息，可以直接调用coco API完成信息的读取。