python TT100K数据集标注转化为yolo格式

一、背景

       由于项目需要进行【交通标志检测】，我们选择TT100K数据集作为我们的训练测试数据集。又因为使用的目标检测网络为YOLO系列，所以需要将TT100K数据集的标注转换成yolo格式的标注，并根据项目需要，将TT100K数据集的交通标志分类分为3类：指示标志（标签为1，对应原数据集“i”开头的类别）、禁令标志（标签为2，对应原数据集“p”开头的类别）、警告标志（标签为3，对应原数据集“w”开头的类别）。

二、TT100K数据集标注文件解析

        TT100K数据集的标注文件为annotation.json文件，打开后发现标注文件是以字典的格式储存标注信息，字典的组织如下所示：

{
"imgs":
    {
    '32773':
	     {'path': 'test/32773.jpg', 
	    'objects':
		     [
		        {'category': 'ph2',
		         'bbox': {'xmin': 924.0, 'ymin': 1132.0, 'ymax': 1177.3333, 'xmax': 966.6667}, 
		        'ellipse_org': [[942.725, 1132.14], [926.19, 1144.18], [922.884, 1156.61], [931.746, 1173.02], [949.471, 1174.07], [959.921, 1169.58], [962.037, 1142.72]],
		         'ellipse': [[943.6439208984375, 1154.060791015625], [41.572391510009766, 45.09551239013672], 36.37429428100586]}, 

		        {'category': 'p11', 
		        'bbox': {'xmin': 970.667, 'ymin': 1128.0, 'ymax': 1170.6667, 'xmax': 1013.3333}, 
		        'ellipse_org': [[997.385, 1130.51], [977.584, 1137.36], [974.222, 1161.52], [990.286, 1170.61], [1012.7, 1157.41]],
		         'ellipse': [[992.5154418945312, 1149.9034423828125], [39.68983840942383, 43.2476692199707], 66.92964172363281]}, 

		        {'category': 'pl5', 
		        'bbox': {'xmin': 1146.67, 'ymin': 1108.0, 'ymax': 1150.6667, 'xmax': 1190.6733000000002}, 
		        'ellipse_org': [[1170.5, 1110.58], [1151.81, 1121.04], [1155.43, 1148.58], [1176.6, 1150.32], [1190.93, 1125.4]], 
		        'ellipse': [[1168.935546875, 1130.9200439453125], [40.40793991088867, 44.34447479248047], 53.05292892456055]}],
	
	     'id': 32773},
    ……
    }
"type": {……}
}

其中imgs保存的是标注信息，type保存的是所有类别的名称。

三、标注转换代码

###############################################################################################
# 这个脚本用于将tt100k数据集转换成yolo格式的标注，用于目标检测网络的训练。
# 其中将交通标志类别分为：指示标志（1）、禁令标志（2）、警告标志（3）
###############################################################################################

import os
import json
import cv2

# def convert(size, box):
#     '''
#     @size: (w, h)， 图片的高宽
#     @box: (xmin, xmax, ymin, ymax), 标注框的坐标
#     @return: (x_center, y_center, w2, h2), 返回目标中心坐标与相对高宽
#     '''
#     dw = 1. / size[0]
#     dh = 1. / size[1]
#     x = (box[0] + box[1]) / 2.0
#     y = (box[2] + box[3]) / 2.0
#     w = box[1] - box[0]
#     h = box[3] - box[2]
#     x = x * dw
#     w = w * dw
#     y = y * dh
#     h = h * dh
#     return (x, y, w, h)


anno_path = r'I:\TT100K\data\annotations.json'
save_path = r'I:\TT100K\data\yolo_annotation'
f = open(anno_path)
anno = json.loads(f.read())
f.close()
classes = anno['types']
imgs = anno['imgs']
# print(imgs.keys())
for key in imgs.keys():
    # print(imgs[key].keys())
    path = imgs[key]['path']
    image = cv2.imread(os.path.join('I:/TT100K/data', path))
    H, W, _ = image.shape
    save_txt_path = os.path.join(save_path, path.replace('/', '\\').replace('.jpg', '.txt'))
    f = open(save_txt_path, 'a')
    for data in imgs[key]['objects']:
        if data['category'].startswith("i"):
            cls = '1'
        elif data['category'].startswith("p"):
            cls = '2'
        else:
            cls = '3'
        x_min = float(data['bbox']['xmin'])
        x_max = float(data['bbox']['xmax'])
        y_min = float(data['bbox']['ymin'])
        y_max = float(data['bbox']['ymax'])
        xc = (x_max + x_min) / 2.0 / W
        yc = (y_max + y_min) / 2.0 / H
        w = (x_max - x_min) / W
        h = (y_max - y_min) / H

        txt = cls + ' ' + str(xc) + ' ' + str(yc) + ' ' + str(w) + ' ' + str(h) + '\n'
        f.writelines(txt)

    f.close()


print(anno.keys())