[YoloV3] DJIＸＭＬ数据集转Ｙｏｌｏ脚本

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'''
Created on Sat Jan 11 16:31:30 2020
XmlToTxt, for DJI ROCO Dataset
@author: HNU robomaster 跃鹿


define:
ArmorBlue=0,
ArmorRed=1,
Base=2,
Watcher=3
Exceptional data will be ignored

'''
import os
import sys
import xml.etree.ElementTree as ET

filedir = 'Q:/RoboMaster/robomaster_North China Regional Competition/image_annotation'
outdir = 'Q:/RoboMaster/robomaster_North China Regional Competition/processedTXT'


def del_all_files(path):
    ls = os.listdir(path)
    for i in ls:
        c_path = os.path.join(path, i)
        if os.path.isdir(c_path):
            del_file(c_path)
        else:
            os.remove(c_path)


def mkdir(path):
    # 去除首位空格
    path = path.strip()
    # 去除尾部 \ 和 / 符号
    path = path.rstrip("\\")
    path = path.rstrip("/")
    # 判断路径是否存在
    is_exist = os.path.exists(path)
    # 判断结果
    if not is_exist:
        os.makedirs(path)

    else:
        # 如果目录存在则首先递归删除该目录下所有内容
        del_all_files(path)


def xml_to_txt(indir, outdir):
    parser = ET.XMLParser(encoding="utf-8")
    root = ET.parse(indir, parser=parser)
    root = root.getroot()
    
    f_w = open(outdir, 'w')
    size_info = root.find("size")
    width = float(size_info.find("width").text)
    height = float(size_info.find("height").text)

    for obj in root.findall("object"):
        #error occurs when using func .find() specifically for string "armor_color"
        color = getattr(obj.find('armor_color'), 'text', None)
        name = obj.find("name").text
        xmin = float(obj.find("bndbox").find("xmin").text)
        ymin = float(obj.find("bndbox").find("ymin").text)
        xmax = float(obj.find("bndbox").find("xmax").text)
        ymax = float(obj.find("bndbox").find("ymax").text)
        x_center = (xmax + xmin) / 2
        y_center = (ymax + ymin) / 2
        x = x_center / width
        y = y_center / height
        w = (xmax - xmin) / width
        h = (ymax - ymin) / height
        flag = 5
        if name == 'armor':
            if color == 'blue':
                flag = 0
            elif color == 'red':
                flag = 1
        elif name == 'base':
            flag = 2
        elif name == 'watcher':
            flag = 3
        if flag != 5:
            f_w.write("".join([str(flag), ' ', str(round(x, 6)), ' ', str(round(y, 6)), ' ', str(round(w, 6)), ' ', str(round(h, 6)), '\n']))
    f_w.close()


def main():
    global filedir
    global outdir

    file_list = os.listdir(filedir)
    mkdir(outdir)
    for file_name in file_list:
        file_prefix = file_name.rpartition(".")[0]
        new_name = "".join([file_prefix, '.txt'])
        xml_to_txt(os.path.join(filedir, file_name), os.path.join(outdir, new_name))


if __name__ == "__main__":
    main()