清华Tinghua100K交通标志数据集转COCO json格式数据集

清华Tinghua100K交通标志数据集转COCO json格式数据集

数据集地址：Tinghua100k
虽然原始标注就已经是json格式了但是数据的排布和COCO的标准不一样，记录下转COCO格式的代码
之前有转过XML了，转XML代码
在XML标记的基础上，转成COCO-json标注：

#!/usr/bin/python

# @ Gyn 6418 2021/10/8

import sys
import os
import json
import xml.etree.ElementTree as ET
import shutil


START_BOUNDING_BOX_ID = 1
PRE_DEFINE_CATEGORIES = {"pr80": 1, "pm50": 2, "ph3.5": 3, "pr20": 4, "p11": 5, "ph2.2": 6, "pr40": 7, "pr10": 8, "i13": 9, "ph5": 10, "p4": 11, "p28": 12, "pa12": 13, "pw4": 14, "pw4.2": 15, "pa13": 16, "pm55": 17, "il110": 18, "pn": 19, "ph4.8": 20, "pm8": 21, "ph4.5": 22, "ph4.3": 23, "pr70": 24, "ph4": 25, "pm35": 26, "pb": 27, "po": 28, "ph3": 29, "ph2.8": 30, "p18": 31, "il100": 32, "pl50": 33, "pw3.2": 34, "pw2": 35, "p22": 36, "p6": 37, "pl30": 38, "ph3.2": 39, "p26": 40, "pm10": 41, "pl5": 42, "i10": 43, "p21": 44, "pm40": 45, "pw3": 46, "pa14": 47, "p15": 48, "p17": 49, "pm20": 50, "il80": 51, "pm30": 52, "pl60": 53, "p16": 54, "ph2": 55, "pl70": 56, "i3": 57, "p5": 58, "i11": 59, "i1": 60, "pl40": 61, "pa10": 62, "pr100": 63, "il50": 64, "pr30": 65, "io": 66, "p19": 67, "il70": 68, "p12": 69, "p13": 70, "pl20": 71, "p3": 72, "i12": 73, "pl100": 74, "pl35": 75, "p25": 76, "pr60": 77, "ph2.9": 78, "ph2.4": 79, "ph2.1": 80, "ph5.3": 81, "i14": 82, "pm2": 83, "ph1.5": 84, "pr50": 85, "p1": 86, "pl80": 87, "pw3.5": 88, "ps": 89, "p27": 90, "il60": 91, "pm15": 92, "ph4.2": 93, "pl25": 94, "p14": 95, "i2": 96, "pne": 97, "pl110": 98, "i4": 99, "pl120": 100, "pl15": 101, "pl10": 102, "p2": 103, "pm13": 104, "pw2.5": 105, "p23": 106, "ph2.5": 107, "p20": 108, "p9": 109, "p10": 110, "i5": 111, "il90": 112, "pl90": 113, "pm5": 114, "p24": 115, "pa8": 116, "pl0": 117, "pw4.5": 118, "p8": 119, "w57": 120, "w55": 121, "wo": 122, "w30": 123, "w35": 124, "w10": 125, "w3": 126, "w18": 127, "w47": 128, "ip": 129, "w38": 130, "w21": 131, "w63": 132, "w45": 133, "w66": 134, "w32": 135, "w42": 136, "w20": 137, "w59": 138, "w16": 139, "w13": 140, "w15": 141, "w5": 142, "w41": 143, "w22": 144, "w46": 145, "w8": 146, "pg": 147, "w34": 148, "w37": 149, "w58": 150, "w12": 151}

def get(root, name):
    vars = root.findall(name)
    return vars

def get_and_check(root, name, length):
    vars = root.findall(name)
    if len(vars) == 0:
        raise NotImplementedError('Can not find %s in %s.'%(name, root.tag))
    if length > 0 and len(vars) != length:
        raise NotImplementedError('The size of %s is supposed to be %d, but is %d.'%(name, length, len(vars)))
    if length == 1:
        vars = vars[0]
    return vars

def get_filename_as_int(filename):
    try:
        filename = os.path.splitext(filename)[0]
        return int(filename)
    except:
        raise NotImplementedError('Filename %s is supposed to be an integer.'%(filename))

def convert(xml_list, xml_dir, json_file):
    allwrongsamples=[]
    list_fp = open(xml_list, 'r')
    json_dict = {"images":[], "type": "instances", "annotations": [],
                 "categories": []}
    categories = PRE_DEFINE_CATEGORIES
    bnd_id = START_BOUNDING_BOX_ID
    for line in list_fp:
        line = line.strip()
        print("Processing %s"%(line))
        xml_f = os.path.join(xml_dir, line)
        tree = ET.parse(xml_f)
        root = tree.getroot()
        path = get(root, 'path')
        if len(path) == 1:
            filename = os.path.basename(path[0].text)
        elif len(path) == 0:
            filename = get_and_check(root, 'filename', 1).text
        else:
            raise NotImplementedError('%d paths found in %s'%(len(path), line))
        ## The filename must be a number
        image_id = get_filename_as_int(filename)
        size = get_and_check(root, 'size', 1)
        width = int(get_and_check(size, 'width', 1).text)
        height = int(get_and_check(size, 'height', 1).text)
        image = {'file_name': filename, 'height': height, 'width': width,
                 'id':image_id}
        json_dict['images'].append(image)
        ## Cruuently we do not support segmentation
        #  segmented = get_and_check(root, 'segmented', 1).text
        #  assert segmented == '0'
        for obj in get(root, 'object'):
            category = get_and_check(obj, 'name', 1).text
            if category not in categories:
                new_id = len(categories)
                categories[category] = new_id
            category_id = categories[category]
            bndbox = get_and_check(obj, 'bndbox', 1)
            xmin = int(get_and_check(bndbox, 'xmin', 1).text)
            ymin = int(get_and_check(bndbox, 'ymin', 1).text)
            xmax = int(get_and_check(bndbox, 'xmax', 1).text)
            ymax = int(get_and_check(bndbox, 'ymax', 1).text)
            try:
                assert(xmax > xmin+1)
                assert(xmin > 1)
                assert(ymin > 1)
                assert(ymax > ymin+1)
                assert(xmax < width-1)
                assert(ymin < height-1)
            except AssertionError:
                allwrongsamples.append(line.split('/')[-1].split('.')[0])
            o_width = abs(xmax - xmin)
            o_height = abs(ymax - ymin)
            ann = {'area': o_width*o_height, 'iscrowd': 0, 'image_id':
                   image_id, 'bbox':[xmin, ymin, o_width, o_height],
                   'category_id': category_id, 'id': bnd_id, 'ignore': 0,
                   'segmentation': []}
            json_dict['annotations'].append(ann)
            bnd_id = bnd_id + 1
    print(allwrongsamples)
    for cate, cid in categories.items():
        cat = {'supercategory': 'none', 'id': cid, 'name': cate}
        json_dict['categories'].append(cat)
    json_fp = open(json_file, 'w')
    json_str = json.dumps(json_dict)
    json_fp.write(json_str)
    json_fp.close()
    list_fp.close()
    return allwrongsamples
    

def removewrong(allwrongsamples):
    imgPath="F:/tt100k_dataset/data/coco/train2017/"
    xmlPath="F:/tt100k_dataset/data/coco/trainXML/"
    for sample in allwrongsamples:
        shutil.move(imgPath+sample+".jpg","F:/tt100k_dataset/data/coco/wrongtrain/"+sample+".jpg")
        shutil.move(xmlPath+sample+".xml","F:/tt100k_dataset/data/coco/wrongtrain/"+sample+".xml")

if __name__ == '__main__':

    allwrongsamples = convert("xml_list_train.txt", "F:/tt100k_dataset/data/coco/trainXML/" ,"instances_train2017.json")
	# removewrong(allwrongsamples)

需要XML目录的txt文件，获取代码：

import os
dir = "F:/tt100k_dataset/data/coco/valXML/"
with open("xml_list_val.txt","a+") as f:
	for root, dirs, files in os.walk(dir):
		for file in files:
			f.write(root+file+"\n")