当前位置：首页 > news >正文

从零搭建pytorch模型教程（八）实践部分（二）目标检测数据集格式转换

news 来源：原创 2024/9/20 7:55:46

前言

图像目标检测领域有一个非常著名的数据集叫做COCO，基本上现在在目标检测领域发论文，COCO是不可能绕过的Benchmark。因此许多的开源目标检测算法框架都会支持解析COCO数据集格式。通过将其他数据集格式转换成COCO格式可以无痛的使用这些开源框架来训练新的数据集，比如Pascal VOC数据集。

本文首先将介绍COCO和VOC目标检测数据集格式以及VOC转换到COCO格式的核心步骤，最后将自定义一种数据格式利用上述核心步骤将其转换到COCO格式下。只要理解了不同数据集的标注方法，转换数据集其实就是一个非常简单自然的过程，可以拓展到任意方式标注的数据集上。

数据集格式介绍

COCO
其实COCO数据集的标签内容不仅仅涵盖目标检测，还包含了目标关键点、实例Mask以及图片描述等信息。在这里我们着重介绍COCO的目标检测相关内容。我们以COCO2017为例先看看其标签文件结构：

其中红框框出来的就是以Json格式组织的目标检测相关的标注文件，其主要由三个部分构成：

"info"字段：数据集的基本信息描述、版本号、年份等信息。

“images字段” ：包含了图片路径、宽高信息、唯一标志ID等信息。

“annotations字段”：包含了图片中的Box位置、类别等信息。

其简单示例如下所示:

{"info": {"description": "COCO 2017 Dataset","url": "http://cocodataset.org","version": "1.0","year": 2017,"contributor": "COCO Consortium","date_created": "2017/09/01"},
"images": [{"license": 4,"file_name": "000000397133.jpg","coco_url": "http://images.cocodataset.org/val2017/000000397133.jpg","height": 427,"width": 640,"date_captured": "2013-11-14 17:02:52","flickr_url": "","id": 397133},{"license": 1,"file_name": "000000037777.jpg","coco_url": "http://images.cocodataset.org/val2017/000000037777.jpg","height": 230,"width": 352,"date_captured": "2013-11-14 20:55:31","flickr_url": "","id": 37777}],"annotations": [{"area": 702.1057499999998,  //Box的尺寸"image_id": 289343,         //对应的图像ID"bbox": [473.07,                 //左上角点x坐标395.93,                 //左上角点y坐标38.65,                  //Box的宽28.67                   //Box的高],"category_id": 18,          //对应的类别"id": 1768,                 //该标签独有ID"iscrowd": 0                //0表示非密集场景，1表示密集场景}]
}

PASCAL VOC

PASCAL VOC数据集有两个相对重要年份的数据集：PASCAL VOC 2007与PASCAL VOC 2012,每年都会在上一年的基础上增加一些额外的数据或标签。PASCAL VOC数据集也涵盖了分类、检测、分割、动作识别等标签。

我们这里着重介绍其检测部分，以PASCAL VOC 2012数据集为例，包含了20个类别1W+数据集，2W+标注Box的目标。其标签格式是每一个图片都有一个对应的XML文件作为其标注信息载体，标注信息主要包含如下几方面内容：

图像基本信息：图像名、图像尺寸等

object字段：目标分类标签、box标签(xmin,ymin,xmax,ymax)等信息

XML主要格式如下：

<annotation><folder>VOC2012</folder><filename>2007_000063.jpg</filename>           //标签对应的图片文件<source><database>The VOC2007 Database</database><annotation>PASCAL VOC2007</annotation><image>flickr</image></source><size>                                         //图像尺寸<width>500</width><height>375</height><depth>3</depth></size><segmented>1</segmented><object><name>dog</name>                           //类别<pose>Unspecified</pose><truncated>0</truncated><difficult>0</difficult>                   //1表示这个目标是比较难识别的<bndbox>                                   //box信息<xmin>123</xmin>                       //左上角x坐标<ymin>115</ymin>                       //左上角y坐标<xmax>379</xmax>                       //右下角x坐标<ymax>275</ymax>                       //右下角y坐标</bndbox></object><object><name>chair</name><pose>Frontal</pose><truncated>1</truncated><difficult>0</difficult><bndbox><xmin>75</xmin><ymin>1</ymin><xmax>428</xmax><ymax>375</ymax></bndbox></object>
</annotation>

数据集格式转换

在知道了各个数据集格式的基础上做数据集格式的转换就已经是非常简单的任务了，也有很多优秀的开源框架已经帮我们做好了这些事情比如MMDetection中就已经提供好了现成的工具供我们白嫖（bushi），使用了。

我们抽取其一些核心部分来一起看看，详细代码请参考MMDetection

Github:https://github.com/openmmlab/mmdetection/tree/master/tools/dataset_converters

从不同的数据集转换到COCO下主要也就两个步骤：

解析待转换数据集格式。

用COCO格式重构Json文件。

上述第二步对任意待转换数据集都是一样的，可以抽象为一个函数，输入的是解析好的不同数据集的Box信息等数据。下面我们以几个不同的数据集为例介绍。

From VOC to COCO
从VOC数据集到COCO数据集格式转换主要包含如下两个步骤：

解析VOC数据集数据:遍历图片以及对应XML文件，返回一个数组A，数组中的每一个实例包含了图片路径、Box相关标注信息等。

遍历A将A中的实例信息用COCO的格式表达出来并生成Json文件

其主要由两块核心代码构成，一个是VOC的XML文件解析，一个是Json文件生成。

VOC XML标注文件解析

xml文件解析已经有下面这个非常方便的Python库供大家使用

import xml.etree.ElementTree as ETdef parse_xml(args):xml_path, img_path = argstree = ET.parse(xml_path)           # 构建XML文件解析树root = tree.getroot()               # 获取XML文件的根节点size = root.find('size')            # 获取图像的尺寸w = int(size.find('width').text)    # 图像宽高h = int(size.find('height').text)bboxes = []labels = []bboxes_ignore = []labels_ignore = []for obj in root.findall('object'):  # 遍历object字段下所有box信息name = obj.find('name').textlabel = label_ids[name]difficult = int(obj.find('difficult').text)  #这个difficult对应的是COCO中iscrowdedbnd_box = obj.find('bndbox')bbox = [                             # 对应的Box标注信息(x1,y1,x2,y2)int(bnd_box.find('xmin').text),int(bnd_box.find('ymin').text),int(bnd_box.find('xmax').text),int(bnd_box.find('ymax').text)]if difficult:                        # 将difficult属性的Box放入ignore列表bboxes_ignore.append(bbox)       # 最后计算AP时这个GT是被忽略的labels_ignore.append(label)else:bboxes.append(bbox)labels.append(label)if not bboxes:bboxes = np.zeros((0, 4))labels = np.zeros((0, ))else:bboxes = np.array(bboxes, ndmin=2) - 1labels = np.array(labels)if not bboxes_ignore:bboxes_ignore = np.zeros((0, 4))labels_ignore = np.zeros((0, ))else:bboxes_ignore = np.array(bboxes_ignore, ndmin=2) - 1labels_ignore = np.array(labels_ignore)annotation = {'filename': img_path,'width': w,'height': h,'ann': {'bboxes': bboxes.astype(np.float32),'labels': labels.astype(np.int64),'bboxes_ignore': bboxes_ignore.astype(np.float32),'labels_ignore': labels_ignore.astype(np.int64)}}return annotation

用解析好的annotation重构COCO格式的Json文件：

import numpy as npdef voc_classes():return ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat','chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person','pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']def cvt_to_coco_json(annotations):image_id = 0annotation_id = 0coco = dict()coco['images'] = []coco['type'] = 'instance'coco['categories'] = []coco['annotations'] = []image_set = set()# coco annotations字段添加标注信息def addAnnItem(annotation_id, image_id, category_id, bbox, difficult_flag):annotation_item = dict()annotation_item['segmentation'] = []# 这里省略了seg部分代码seg = []annotation_item['segmentation'].append(seg)# 转换为COCO对应的x1,y1,w,h格式xywh = np.array([bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]])annotation_item['area'] = int(xywh[2] * xywh[3])# 如果difficult标志为1，该目标对应coco中iscrowd字段为1表明为密集目标场景if difficult_flag == 1:annotation_item['ignore'] = 0annotation_item['iscrowd'] = 1else:annotation_item['ignore'] = 0annotation_item['iscrowd'] = 0annotation_item['image_id'] = int(image_id)annotation_item['bbox'] = xywh.astype(int).tolist()annotation_item['category_id'] = int(category_id)annotation_item['id'] = int(annotation_id)coco['annotations'].append(annotation_item)return annotation_id + 1#for category_id, name in enumerate(voc_classes()):category_item = dict()category_item['supercategory'] = str('none')category_item['id'] = int(category_id)category_item['name'] = str(name)coco['categories'].append(category_item)for ann_dict in annotations:file_name = ann_dict['filename']ann = ann_dict['ann']assert file_name not in image_setimage_item = dict()image_item['id'] = int(image_id)image_item['file_name'] = str(file_name)image_item['height'] = int(ann_dict['height'])image_item['width'] = int(ann_dict['width'])coco['images'].append(image_item)        # 设置COCO的"images"字段image_set.add(file_name)# 设置COCO的"annotations"字段bboxes = ann['bboxes'][:, :4]            # 获取box和label类别信息labels = ann['labels']for bbox_id in range(len(bboxes)):bbox = bboxes[bbox_id]label = labels[bbox_id]annotation_id = addAnnItem(annotation_id, image_id, label, bbox, difficult_flag=0)# ignore的目标表示该GT被忽视bboxes_ignore = ann['bboxes_ignore'][:, :4]labels_ignore = ann['labels_ignore']for bbox_id in range(len(bboxes_ignore)):bbox = bboxes_ignore[bbox_id]label = labels_ignore[bbox_id]annotation_id = addAnnItem(annotation_id, image_id, label, bbox, difficult_flag=1)image_id += 1return coco

拿到返回的coco对象后只需要调用下列方法就可以将对象序列化成Json文件了。

import mmcv
mmcv.dump(coco, out_file) # out_file为输出的json文件名

值得注意的一点是上面提到的iscorwd这个字段，这个字段标注为1时，最后统计AP时，该GT与预测框完成匹配后还可以考虑与其他预测框进行匹配，允许多个预测框与其匹配(因为场景是密集的)。

自定义格式数据集 to COCO

首先我们自定义一种数据标注格式，我们用txt文件作为标注信息的载体，将txt文件与图像文件通过相同的文件名一一对应。分别将标签文件以及对应图像文件放在Annotations以及JPEGImages文件夹下，同时我们生成JPEGImages图像文件的filelist.txt文件，这个文件每一行对应一个图像文件的全路径：

txt文件格式如下：

第一列表示类别，从0开始；第二到第五列表示Box信息依次为中心点x方向坐标，中心点y方向坐标，box的宽以及高(cx,cy,w,h)。

我们同样使用前面介绍过的cvt_to_coco_json将固定格式的annotations转换为COCO格式，那么我们只需要编写解析自定义格式数据集生成annotations的代码即可：

box尺寸小于min_size的作为ignore对象

file_path为图像路径的filelist.txt文件的全路径

def parse_info(file_path, min_size):annotations = []invalid_img = 0small_box = 0with open(file_path, 'r') as f:for l in tqdm(f):img_file = l.rstrip()img = cv2.imread(img_file)if img is None:invalid_img += 1continueh,w,_ = img.shape# 获取对应的标签文件ann_file = img_file.replace("JPEGImages", "Annotations").replace \(".png", ".txt").replace(".jpg", ".txt")annotation = {'filename' : img_file,'height' : h,'width' : w,'ann' : {}}if not osp.exists(ann_file):annotations.append(annotation)continueboxes, labels = [], []boxes_ignore, labels_ignore = [], []with open(ann_file, 'r') as fr:for anno in fr:anno_list = anno.rstrip().split(' ')cls = int(anno_list[0])cx, cy = float(anno_list[1]), float(anno_list[2])w, h = float(anno_list[3]), float(anno_list[4])# 转换为COCO box表示格式x1 = max(0, int(cx-w/2))y1 = max(0, int(cy-h/2))box = [x1, y1, w, h]if w >= min_size and h >= min_size:labels.append(cls)boxes.append(box)else:labels_ignore.append(cls)boxes_ignore.append(box)                    boxes = np.zeros((0, 4)) if len(boxes) == 0 else np.array(boxes)labels = np.zeros((0, )) if len(labels) == 0 else np.array(labels)      boxes_ignore = np.zeros((0, 4)) if len(boxes_ignore) == 0 else np.array(boxes_ignore)labels_ignore = np.zeros((0, )) if len(labels_ignore) == 0 else np.array(labels_ignore)  annotation['ann']['bboxes'] = np.array(boxes)annotation['ann']['labels'] = np.array(labels)annotation['ann']['bboxes_ignore'] = np.array(boxes_ignore)annotation['ann']['labels_ignore'] = np.array(labels_ignore)annotations.append(annotation)print('INFO：Invalid IMG:{}'.format(invalid_img))return annotations