#!/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()