COCO 数据集格式及mmdetection中的转换方法

COCO 数据集格式及mmdetection中的转换方法

COCO格式

CV中的目标检测任务不同于分类，其标签的形式稍为复杂，有几种常用检测数据集格式，本文将简要介绍最为常见的COCO数据集的格式。

完整的官方样例可自行查阅，以下是几项关键的字段：

{
    "images": [image],
    "annotations": [annotation],
    "categories": [category]
}


image = {
    "id": int,
    "width": int,
    "height": int,
    "file_name": 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,
}]

以下是一组实例：

dict{
  'images': list[
    	dict{
        'id': 0,
        'file_name': 34020010494_e5cb88e1c4_k.jpg,
        'height': 1536,
        'width': 2048
      }, ...
    'annotations': list[
      	dict{
          'image_id': 0,
          'id': 1,
          'category_id': 0,
          'bbox': [135, 115, 676, 850],
          'area': 574600,
          'segmentation': [[586.5, ...], []],
          'is_crowd': 0
        }, ...
      'catgories': list[
        	dict{
            'id': 0,
            'name': 'balloon'
          },...
      ]
    ]
  ]
}

首先最外层（json文件读进来）是一个字典，该字典有三个键：'images', 'annotations', 'categories'；

​		其中'images'是一个列表，长度就是数据集图像数；

​				列表中每个元素又是一个字典，有四个键：'id', 'file_name', 'height', 'width'，和他们的值

​		其中'annotations'也是一个列表，长度就是数据集所有标注（instances）数；

​				列表中每个元素又是一个字典，有七个键：'image_id', 'id', 'category_id', 'bbox', 'area', 'segmentation', 'is_crowd'

​		其中'categories'是一个列表，长度为数据集所含有的类别数；

​				列表中每个元素又是一个字典，有两个键：'id', 'name'

mmdetection中的转换方法

官方例程

mmdetection中提供了几种训练自己的数据集的方法，其中之一就是将自己的数据集转换为COCO格式，进行训练。

官方给出了balloon数据集的转换样例，初次尝试可以按照此例程学习，这里同时提供balloon数据集下载链接。

其他例程

除此之外，笔者想在这里介绍一下另一个非官方的数据集，转换为COCO格式的例程，供读者参考。

首先介绍一下这个pig数据集。

与balloon数据集类似，pig数据集只有pig一个类，共700个样本，原始给出的数据集格式如下：

1. 文件路径

   images
   	xxx.jpg
   	...
   labels
   	xxx.json
   	...
   

   本数据集有两个文件夹，分别存放images和labels，对应文件名相同，扩展名分别为jpg和json。

2. 标注格式

   json文件中给出的标注格式如下：

   dict{
     'shape': list[
       	dict{
           'label': 'pig',
           'boxes': [x1, x2, y1, y2],
           'points': [[x1, y1], [x2, y2], ..., [xn, yn]]
   			}
       'imagePath': 'xxx.jpg'
     ]
   }
   

   json文件最外层是一个字典，该字典有两个键：'shape'，'imagePath';
   
   ​	其中'shape'内是一个列表，列表的长度是该张图片内pig的个数；
   
   ​		列表内每个元素又是一个字典，有三个键：'label'，'boxes'，'points'；
   
   ​	其中'imagePath'内则是本标注文件所对应的图片名。
   

   原本数据集内700个样本有500个box标注，其他200个是mask标注，经过我的初步处理，已经将没有box标注的样本根据mask标注生成了box标注。其中box标注是给出了左上右下两个对角坐标，而mask标注，也是给出了全部点的坐标。注意此处与COCO数据集给出标注的格式不同，需要做一定的转换。

转换为COCO数据集格式的函数如下：

import os
import os.path as osp
import mmcv

def convert_pig_to_coco(img_dir, label_dir, out_file):
    samples_list = [sample.split('.')[0] for sample in os.listdir(label_dir)]

    obj_count = 0
    images = []
    annotations = []
    for image_id, sample in enumerate(mmcv.track_iter_progress(samples_list)):
        prev_dict = mmcv.load(osp.join(label_dir, sample+'.json'))
        height, width = mmcv.imread(osp.join(img_dir, sample+'.jpg')).shape[: 2]
        images.append(dict(
            id = image_id,
            file_name = sample+'.jpg',
            height = height,
            width = width))
        for pig in prev_dict['shape']:
            # print(pig['boxes'])
            x_min, y_min, x_max, y_max = pig['boxes']
            mask = pig['points']
            if mask != None:
                px = [item[0] for item in mask]
                py = [item[1] for item in mask]
                poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
                poly = [p for x in poly for p in x]
            else:
                poly = None
            
        
            data_anno = dict(
                    image_id=image_id,
                    id=obj_count,
                    category_id=0,
                    bbox=[x_min, y_min, x_max - x_min, y_max - y_min],
                    area=(x_max - x_min) * (y_max - y_min),
                    segmentation=[poly],
                    iscrowd=0)
            annotations.append(data_anno)
            obj_count += 1
    coco_format_json = dict(
                        images=images,
                        annotations=annotations,
                        categories=[{'id':0, 'name': 'pig'}])
    mmcv.dump(coco_format_json, out_file)
        

提供这个额外的例程是为了说明不必拘泥于官方给出的转换例程框架，要抓住数据集的本质，按照合理的方式将原格式内的数据读取出来并以COCO数据集的格式存放到一整个json文件中即可。只要完成这项任务，就是成功完成了自己的数据集到COCO格式的数据集的转换。
另外，如果有与本数据集原格式高度相似的检测数据集要转换为COCO格式，也可直接参考本例程的内容。