需要修改的部分:
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改成自己大类下面的细分属性
转换代码如下, 此处的代码的只把 traffic light 类转换, 不区别red, green, yellow等:
import os
import os.path as osp
import json
import shutil
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
# bdd100k数据的根目录
BDD_FOLDER = "/home/wsy/data/06_dataset_transform/01_trafficlight/bdd100k_kaggle"
# 保存转换后筛选出的VOC的图片目录
img_dir = "/home/wsy/data/06_dataset_transform/01_trafficlight/bdd100ktovoc2007/images"
if DEBUG:
XML_PATH = "./xml"
else:
# 保存转换后筛选出的VOC的标签目录
XML_PATH = "/home/wsy/data/06_dataset_transform/01_trafficlight/bdd100ktovoc2007/xml"
def mkr(path):
if os.path.exists(path):
shutil.rmtree(path)
os.mkdir(path)
else:
os.makedirs(path, exist_ok=True) # os.makedirs() 递归创建文件夹
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)
mkr(osp.join(img_dir, 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))
# save selects image to another folder.
img_path = osp.join(image_folder, osp.splitext(datum['name'])[0] + '.jpg')
dst_imgpath = osp.join(img_dir, trainval)
shutil.copy(img_path, dst_imgpath)
print("all of classes are ", 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)
成功运行后会生成:
XML_PATH目录,在该目录下有,xml文件的train和val两个目录,里面存放的为VOC数据集所需xml注释文件。
img_dir目录下, 会有筛选出的与xml文件对应的train和val图片.
借鉴文章: https://blog.csdn.net/l297969586/article/details/89248108