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

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

COCO数据集可视化

- 1.标注框bbox的单独可视化
- - 1.1 利用opencv可视化
  - 1.2 利用coco API读取json
- 2.分割标注segmentation的单独可视化
- 3.bbox和segmentation的联合可视化

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

1.1 利用opencv可视化

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

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文件中的类别，图片，标注信息。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

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

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

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同时标注出来。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

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完成信息的读取。