TT100K2021数据集标注格式转换

TT100K数据集简介：
TT100K数据集是一个公开的交通标志数据集，由清华大学和腾讯联合发布。下载地址(Tsinghua-Tencent 100K Annotations 2021  (with more classification))：http://cg.cs.tsinghua.edu.cn/traffic-sign/数据集大概17.8G，包含4个文件夹，其中三个文件夹，分别存放的图片是 train文件夹存放了 6107张图片，test 文件夹存放了 3073 张图片，other 存放了7643 张（没标签 没sign的）图片。annotations_all.json存放的是数据集的标注信息。

数据集格式转换 

TT100K 数据集的标注文件以 JSON 格式提供，描述了每张图片中交通标志的边界框坐标和类别标签。例如：

{
  "types": ["pl80", "w9", ...],
  "imgs":{
    "62627": {
      "path": "train/62627.jpg", 
	  "id": 62627, 
	  "objects": [
	    {
	      "bbox": {"xmin": 1580.0, "ymin": 758.667, "xmax": 1638.6667, "ymax": 818.6667}, 
		  "category": "ph5"
	    },
		...
	  ]
    }
	...
}

 转为VOC(标注文件xml结尾)

完整转换代码：

import os
import cv2
import json
import xml.dom.minidom


class TT100K2VOC:
    def __init__(self):
        self.dataset_path: str = 'D:\\TT100K2021'  # 修改为自已存放TT100K2021的目录
        self.json_file: str = 'annotations_all.json'
        self.xml_path: str = 'D:\\TT100K2021\\Annotations'  # VOC标注保存目录

    @staticmethod
    def xml_create_annotations(doc: xml.dom.minidom.Document, image_path: str, dsize: tuple):
        annotation = doc.createElement('annotation')
        doc.appendChild(annotation)

        folder = doc.createElement('folder')
        folder_text = doc.createTextNode('JPEGImages')
        folder.appendChild(folder_text)
        annotation.appendChild(folder)

        filename = doc.createElement('filename')
        filename_text = doc.createTextNode(os.path.basename(image_path))
        filename.appendChild(filename_text)
        annotation.appendChild(filename)

        path = doc.createElement('path')
        path_text = doc.createTextNode(image_path.replace('/', '\\'))
        path.appendChild(path_text)
        annotation.appendChild(path)

        source = doc.createElement('source')
        database = doc.createElement('database')
        database_text = doc.createTextNode('Unknown')
        source.appendChild(database)
        database.appendChild(database_text)
        annotation.appendChild(source)

        size = doc.createElement('size')
        width = doc.createElement('width')
        size.appendChild(width)
        width_text = doc.createTextNode(str(dsize[1]))
        width.appendChild(width_text)
        height = doc.createElement('height')
        size.appendChild(height)
        height_text = doc.createTextNode(str(dsize[0]))
        height.appendChild(height_text)
        depth = doc.createElement('depth')
        size.appendChild(depth)
        depth_text = doc.createTextNode('3')
        depth.appendChild(depth_text)
        annotation.appendChild(size)

        segmented = doc.createElement('segmented')
        segmented_text = doc.createTextNode('0')
        segmented.appendChild(segmented_text)
        annotation.appendChild(segmented)
        return annotation

    @staticmethod
    def xml_create_object(doc: xml.dom.minidom.Document, category: str, bbox: tuple):
        objectNode = doc.createElement('object')

        name = doc.createElement('name')
        objectNode.appendChild(name)
        name_text = doc.createTextNode(category)
        name.appendChild(name_text)

        pose = doc.createElement('pose')
        objectNode.appendChild(pose)
        pose_text = doc.createTextNode('Unspecified')
        pose.appendChild(pose_text)

        truncated = doc.createElement('truncated')
        objectNode.appendChild(truncated)
        truncated_text = doc.createTextNode('0')
        truncated.appendChild(truncated_text)

        difficult = doc.createElement('difficult')
        objectNode.appendChild(difficult)
        difficult_text = doc.createTextNode('0')
        difficult.appendChild(difficult_text)

        box = doc.createElement('bndbox')
        objectNode.appendChild(box)
        x_min = doc.createElement('xmin')
        box.appendChild(x_min)
        x_min_text = doc.createTextNode(str(bbox[0]))
        x_min.appendChild(x_min_text)
        y_min = doc.createElement('ymin')
        box.appendChild(y_min)
        y_min_text = doc.createTextNode(str(bbox[1]))
        y_min.appendChild(y_min_text)
        x_max = doc.createElement('xmax')
        box.appendChild(x_max)
        x_max_text = doc.createTextNode(str(bbox[2]))
        x_max.appendChild(x_max_text)
        y_max = doc.createElement('ymax')
        box.appendChild(y_max)
        y_max_text = doc.createTextNode(str(bbox[3]))
        y_max.appendChild(y_max_text)
        return objectNode

    def json2xml(self):
        json_data = json.loads(open(self.dataset_path + "/" + self.json_file).read())
        images = json_data['imgs']

        for index, image_id in enumerate(images):
            image_data = images[image_id]
            save_xml_dir = os.path.join(self.xml_path, os.path.dirname(image_data['path']))
            os.makedirs(save_xml_dir, exist_ok=True)

            image_path = self.dataset_path + '/' + image_data['path']
            if not os.path.exists(image_path):
                continue
            image = cv2.imread(image_path)

            doc = xml.dom.minidom.Document()
            annotation = self.xml_create_annotations(doc, image_path, image.shape[:2])

            image_objects = image_data['objects']
            for item in image_objects:
                object_category = item['category']
                object_bbox = item['bbox']
                x_min, y_min = int(object_bbox['xmin']), int(object_bbox['ymin'])
                x_max, y_max = int(object_bbox['xmax']), int(object_bbox['ymax'])
                objectNode = self.xml_create_object(doc, object_category, (x_min, y_min, x_max, y_max))
                annotation.appendChild(objectNode)
            print(f'正在处理第{index + 1}个文件：{image_path}')
            xml_file = open(save_xml_dir + '/%s.xml' % image_id, 'w+')
            doc.writexml(xml_file, addindent='\t', newl='\n', encoding='utf-8')
            xml_file.close()


if __name__ == '__main__':
    tt100k = TT100K2VOC()
    tt100k.json2xml()

转为YOLO(标注文件txt结尾)

完整转换代码：

import cv2
import os
import json


class TT100K2YOLO:
    def __init__(self):
        self.dataset_path: str = 'D:\\TT100K2021'  # 下载的TT100K数据集所在目录
        self.json_file: str = 'annotations_all.json'  # TT100K中的标注信息
        self.save_path: str = 'D:\\TT100K2021\\YOLO'  # 转换后的YOLO标注保存目录

    @staticmethod
    def convert(size, box):
        dw = 1. / (size[0])
        dh = 1. / (size[1])
        x = (box[0] + box[2]) / 2.0
        y = (box[1] + box[3]) / 2.0
        w = box[2] - box[0]
        h = box[3] - box[1]
        # round函数确定(xmin, ymin, xmax, ymax)的小数位数
        x = round(x * dw, 6)
        w = round(w * dw, 6)
        y = round(y * dh, 6)
        h = round(h * dh, 6)
        return x, y, w, h

    def json2yolo(self):
        # 读TT100K原始数据集标注文件
        with open(os.path.join(self.dataset_path, self.json_file)) as data_json:
            data_dict = json.load(data_json)
            classes = list(data_dict['types'])
            images = data_dict['imgs']

        for imageId in images:
            objects = images[imageId]['objects']
            if len(objects) == 0:
                continue

            path = images[imageId]['path']
            image_path = self.dataset_path + '/' + path
            if not os.path.exists(image_path):
                continue

            image = cv2.imread(image_path)
            save_dir = os.path.join(self.save_path, os.path.dirname(path))
            os.makedirs(save_dir, exist_ok=True)
            with open(os.path.join(save_dir, imageId + '.txt'), 'w') as txt_file:
                for item in objects:
                    bbox, category = item['bbox'], item['category']
                    box = (bbox['xmin'], bbox['ymin'], bbox['xmax'], bbox['ymax'])
                    x, y, w, h = self.convert(image.shape[:2], box)
                    txt_file.write(f'{classes.index(category)} {x} {y} {w} {h}' + '\n')
            print(f'生成文件：{save_dir + "/" + imageId + ".txt"}')

        with open(os.path.join(self.save_path, 'classes.txt'), 'w') as file:
            for line in classes:
                file.write(line + '\n')


if __name__ == '__main__':
    tt100k = TT100K2YOLO()
    tt100k.json2yolo