COCO (Common Objects in COntext) 是一个大型的图像数据集,提供了目标检测、分割、看图说话等多个任务的标签。COCO的标注文件是用json格式编写的,初次接触时需要花十来分钟熟悉一下COCO的标注格式。
本文将简明地介绍COCO目标检测数据集的读取方法,并给出可以调用的Python脚本。读取其他任务的标签时也可以借鉴这些思路。
完整代码:https://github.com/SingleZombie/DL-Demos/tree/master/dldemos/MyYOLO/load_coco.py
COCO的官方网站给出了标注格式的介绍。其中,目标检测的格式是这样的:
{
"info": info,
"images": [image],
"annotations": [annotation],
"licenses": [license],
}
info{
"year": int,
"version": str,
"description": str,
"contributor": str,
"url": str,
"date_created": datetime,
}
image{
"id": int,
"width": int,
"height": int,
"file_name": str,
"license": int,
"flickr_url": str,
"coco_url": str,
"date_captured": datetime,
}
license{
"id": int,
"name": str,
"url": str,
}
annotation{
"id": int,
"image_id": int,
"category_id": int,
"segmentation": RLE or [polygon],
"area": float,
"bbox": [x,y,width,height],
"iscrowd": 0 or 1,
}
categories[{
"id": int,
"name": str,
"supercategory": str,
}]
这里面很多信息是冗余的。我们主要关心图像的信息和检测框的信息。假设该json文件读取进来后叫做root
,我们就可以用root['images']
获取图像信息的列表。在列表中,每一条图像信息的主要属性有:
image{
"id": int,
"width": int,
"height": int,
"file_name": str,
}
id
是用来唯一标记一张图片的。之后我们需要根据这个id
把图像与其标签绑定起来。其他三个属性都是常见的图像属性。
我们可以用root['annotation']
获取标注信息的列表。每一条标注信息表示一个物体的检测信息,一幅图可能有多条检测信息。在目标检测中,我们主要关注检测框位置、大小、类别这几个信息。因此,我们最终要关注的属性有:
annotation{
"id": int,
"image_id": int,
"category_id": int,
"bbox": [x,y,width,height],
}
image_id
与之前图像的id
对应。我们可以根据这个域把图像和标注绑定起来。category_id
可以与下文介绍的分类类别信息绑定起来。bbox
则标记出了每一个检测框的位置和大小。’
最后,我们要获取每个类别id
对应的类别名,方便检验目标检测的结果是否正确。类别信息的列表可以由root['categories']
获得,每一条记录的格式为:
categories[{
"id": int,
"name": str,
"supercategory": str,
}]
其中,id
与前文categroy_id
对应。name
是具体的类别,supercategory
是大类。一般我们只关注name
就行。
知道了数据格式的原理,就可以写脚本读取它们了。
首先介绍一下Python的json库的用法。使用该库时,要先import json
,之后用json.load(fp)
就可以读取一个被打开的json文件指针fp
了。以下是一个示例,路径请根据实际情况自行更改。
import json
def print_json():
with open('data/coco/annotations/instances_val2014.json') as fp:
root = json.load(fp)
打开了文件后,我们可以输出一些信息,熟悉一下json的API,同时具体查看一下COCO标注文件的格式。
import json
def print_json():
with open('data/coco/annotations/instances_val2014.json') as fp:
root = json.load(fp)
print('info:')
print(root['info'])
print('categories:')
print(root['categories'])
print('Length of images:', len(root['images']))
print(root['images'][0])
print('Length of annotations:', len(root['annotations']))
print(root['annotations'][0])
def main():
print_json()
if __name__ == '__main__':
main()
json库能以Python词典的形式访问json的对象,以列表的形式访问json的数组。这段代码中,root
是文件的根节点。root['info']
就是根节点的info
属性的对象。root['categories']
, root['images']
, root['annotations']
分别是类别、图像、标注信息的列表。
运行这个脚本,我的输出大概是:
info:
{'description': 'COCO 2014 Dataset', 'url': 'http://cocodataset.org', 'version': '1.0', 'year': 2014, 'contributor': 'COCO Consortium', 'date_created': '2017/09/01'}
categories:
[{'supercategory': 'person', 'id': 1, 'name': 'person'} ...
Length of images: 40504
{'license': 3, 'file_name': 'COCO_val2014_000000391895.jpg', 'coco_url': 'http://images.cocodataset.org/val2014/COCO_val2014_000000391895.jpg', 'height': 360, 'width': 640, 'date_captured': '2013-11-14 11:18:45', 'flickr_url': 'http://farm9.staticflickr.com/8186/8119368305_4e622c8349_z.jpg', 'id': 391895}
Length of annotations: 291875
{'segmentation': [[239.97, 260.24, 222.04, 270.49, 199.84, 253.41, 213.5, 227.79, 259.62, 200.46, 274.13, 202.17, 277.55, 210.71, 249.37, 253.41, 237.41, 264.51, 242.54, 261.95, 228.87, 271.34]], 'area': 2765.1486500000005, 'iscrowd': 0, 'image_id': 558840, 'bbox': [199.84, 200.46, 77.71, 70.88], 'category_id': 58, 'id': 156}
接下来,我们可以编写一个读取图像文件名及其对应检测框的函数。利用这个函数,我们就可以构建一个目标检测项目的数据集了。
def load_img_ann(ann_path='data/coco/annotations/instances_val2014.json'):
"""return [{img_name, [{x, y, h, w, label}]}]"""
with open(ann_path) as fp:
root = json.load(fp)
img_dict = {}
for img_info in root['images']:
img_dict[img_info['id']] = {'name': img_info['file_name'], 'anns': []}
for ann_info in root['annotations']:
img_dict[ann_info['image_id']]['anns'].append(
ann_info['bbox'] + [ann_info['category_id']])
return img_dict
在这个函数中,我们想构造一个词典img_dict
。它的key是图像id,value是一个属性词典。属性词典的格式是:
{
'name': ...,
'anns': [[x, y, w, h, label], ...]
}
图像文件名name
可以从root['images']
里获取,标注信息可以从root['annotations']
里获取。跑两个循环,取出对应的信息,把信息组合一下塞入词典即可。
注意! COCO 2014里的category id不是连续的。看上去category id最多到了90,但实际上一共只有80个类别。在项目中,还应该自己写一层0-79到categroy id的映射。
在自己的目标检测项目中,直接调load_img_ann(ann_path)
就行了。根据具体的目标检测算法,再进一步预处理检测框。
我自己的项目里有一个可视化bbox的函数draw_bbox
。为了验证读取数据集的函数是否正确,我还写了一个可视化COCO标签的函数。整个脚本如下:
import json
import os
def load_img_ann(ann_path='data/coco/annotations/instances_val2014.json'):
"""return [{img_name, [ (x, y, h, w, label), ... ]}]"""
with open(ann_path) as fp:
root = json.load(fp)
img_dict = {}
for img_info in root['images']:
img_dict[img_info['id']] = {'name': img_info['file_name'], 'anns': []}
for ann_info in root['annotations']:
img_dict[ann_info['image_id']]['anns'].append(
ann_info['bbox'] + [ann_info['category_id']])
return img_dict
def show_img_ann(img_info):
from PIL import Image
from dldemos.nms.show_bbox import draw_bbox
print(img_info)
with open('data/coco/annotations/instances_val2014.json') as fp:
root = json.load(fp)
categories = root['categories']
category_dict = {int(c['id']): c['name'] for c in categories}
img_path = os.path.join('data/coco/val2014', img_info['name'])
img = Image.open(img_path)
for ann in img_info['anns']:
x, y, w, h = ann[0:4]
x1, y1, x2, y2 = x, y, x + w, y + h
draw_bbox(img, (x1, y1, x2, y2), 1.0, text=category_dict[ann[4]])
img.save('work_dirs/tmp.jpg')
def main():
img_dict = load_img_ann()
keys = list(img_dict.keys())
show_img_ann(img_dict[keys[1]])
if __name__ == '__main__':
main()
里面的所有路径都需要根据实际情况进行修改。
这份脚本的可视化结果work_dirs/tmp.jpg
如下:
可以看出,读取出来的标签都是对的。我写的这个API可以放心地调用。