无人驾驶bdd数据集转voc格式

淩亂°似流年 2022-02-20 06:06 307阅读 0赞

需要修改的部分：

1、BDD\_FOLDER：修改成自己的bdd数据集root路径

2、由于我训练的为traffic light类，且类别为\[‘red’,‘green’,‘yellow’,‘none’\]，这些属性属于  
label\[‘attributes’\]\[“trafficLightColor”\]，而不属于大类label\[‘category’\]  
因此，如果你要训练大类，例如训练’traffic light’类，而不是细分的\[‘red’,‘green’,‘yellow’,‘none’\]，则需修改

59行把traffic light改为自己要训练的类
    注释63行
    解注释72行
    注释73行
    解注释77行
    注释78行

3、训练细分类

59行把traffic light改为自己要训练的类
    63行把trafficLightColor改成自己大类下面的细分属性

**转换代码如下：**

import os
    import os.path as osp
    
    import json
    
    from xml.etree.ElementTree import Element, SubElement
    from xml.etree import ElementTree
    from xml.dom import minidom
    
    from PIL import Image
    
    from tqdm import tqdm
    
    DEBUG = False
    
    BDD_FOLDER = "./bdd100k"
    
    if DEBUG:
        XML_PATH = "./xml"
    else:
        XML_PATH = BDD_FOLDER + "/xml"
    
    
    def bdd_to_voc(bdd_folder, xml_folder):
        image_path = bdd_folder + "/images/100k/%s"
        label_path = bdd_folder + "/labels/bdd100k_labels_images_%s.json"
    
        classes = set()
    
        for trainval in ['train', 'val']:
            image_folder = image_path % trainval
            json_path = label_path % trainval
            xml_folder_ = osp.join(xml_folder, trainval)
    
            if not os.path.exists(xml_folder_):
                os.makedirs(xml_folder_)
    
            with open(json_path) as f:
                j = f.read()
            data = json.loads(j)
            for datum in tqdm(data):
                tmp_list = []
                annotation = Element('annotation')
                SubElement(annotation, 'folder').text ='VOC2007'
                SubElement(annotation, 'filename').text = datum['name']
    	    source = get_source()
    	    owner = get_owner()
    	    annotation.append(source)
    	    annotation.append(owner)
                size = get_size(osp.join(image_folder, datum['name']))
                annotation.append(size)
                SubElement(annotation, 'segmented').text ='0'
                # additional information
                #for key, item in datum['attributes'].items():
                #    SubElement(annotation, key).text = item
    
                # bounding box
                for label in datum['labels']:
                    if label['category'] != "traffic light":
                        continue
                    else:
                        tmp_list.append(1)
                    color = label['attributes']["trafficLightColor"]
                    try:
                        box2d = label['box2d']
                    except KeyError:
                        continue
                    else:
                        bndbox = get_bbox(box2d)
    
                    object_ = Element('object')
                    #SubElement(object_, 'name').text = label['category']
                    SubElement(object_, 'name').text = color
                    SubElement(object_, 'pose').text = "Unspecified"
                    SubElement(object_, 'truncated').text = '0'
                    SubElement(object_, 'difficult').text = '0'
                    #classes.add(label['category'])
                    classes.add(color)
    
                    object_.append(bndbox)
                    annotation.append(object_)
                if len(tmp_list) == 0:
                    continue
                xml_filename = osp.splitext(datum['name'])[0] + '.xml'
                with open(osp.join(xml_folder_, xml_filename), 'w') as f:
                    f.write(prettify(annotation))
        print(classes)
    
    def get_owner():
        owner = Element('owner')
        SubElement(owner, 'flickrid').text ='NULL'
        SubElement(owner, 'name').text ='lijing'
        return owner
    
    def get_source():
        source = Element('source')
        SubElement(source, 'database').text ='voc_bdd'
        SubElement(source, 'annotation').text ='VOC2007'
        SubElement(source, 'image').text ='flickr'
        SubElement(source, 'flickrid').text ='NULL'
        return source
    
    
    
    
    def get_size(image_path):
        i = Image.open(image_path)
        sz = Element('size')
        SubElement(sz, 'width').text = str(i.width)
        SubElement(sz, 'height').text = str(i.height)
        SubElement(sz, 'depth').text = str(3)
        return sz
    
    
    def get_bbox(box2d):
        bndbox = Element('bndbox')
        SubElement(bndbox, 'xmin').text = str(int(round(box2d['x1'])))
        SubElement(bndbox, 'ymin').text = str(int(round(box2d['y1'])))
        SubElement(bndbox, 'xmax').text = str(int(round(box2d['x2'])))
        SubElement(bndbox, 'ymax').text = str(int(round(box2d['y2'])))
        return bndbox
    
    
    def prettify(elem):
        rough_string = ElementTree.tostring(elem, 'utf-8')
        reparsed = minidom.parseString(rough_string)
        return reparsed.toprettyxml(indent="\t")
    
    
    if __name__ == "__main__":
        bdd_to_voc(BDD_FOLDER, XML_PATH)

成功运行后会生成  
BDD\_FOLDER/xml目录，在该目录下有train和val两个目录，里面存放的为VOC数据集所需xml注释文件。