diff --git a/convert-annotations.py b/convert-annotations.py new file mode 100644 index 0000000..97da64a --- /dev/null +++ b/convert-annotations.py @@ -0,0 +1,56 @@ +import os +import shutil +import xml.etree.cElementTree as et +from pathlib import Path + +labels_dir = 'labels' +images_dir = 'images' + +tag_size = 'size' +tag_object = 'object' +tag_name = 'name' +tag_bndbox = 'bndbox' +# 1 Вантажівка, 2 Машина легкова +name_class_map = {'Truck': 1, 'Car': 2, 'Taxi': 2} + + +def convert_xml(folder): + for f in os.listdir(folder): + if not f.endswith('.jpg'): + continue + + os.makedirs(images_dir, exist_ok=True) + os.makedirs(labels_dir, exist_ok=True) + + shutil.copy(os.path.join(folder, f), os.path.join(images_dir, f)) + label = f'{Path(f).stem}.xml' + + tree = et.parse(os.path.join(folder, label)) + root = tree.getroot() + lines = [] + size_dict = {size_ch.tag : size_ch.text for size_ch in root.findall(f'{tag_size}/*')} + width = int(size_dict['width']) + height = int(size_dict['height']) + for node_object in tree.findall(tag_object): + for node_object_ch in node_object: + if node_object_ch.tag == tag_name: + class_num = name_class_map[node_object_ch.text] + if node_object_ch.tag == tag_bndbox: + bbox_dict = {bbox_ch.tag: bbox_ch.text for bbox_ch in node_object_ch} + xmin = int(bbox_dict['xmin']) + xmax = int(bbox_dict['xmax']) + ymin = int(bbox_dict['ymin']) + ymax = int(bbox_dict['ymax']) + c_w = (xmax - xmin) / width + c_h = (ymax - ymin) / height + c_x = xmin / width + c_w / 2 + c_y = ymin / height + c_h / 2 + lines.append(f'{class_num} {c_x} {c_y} {c_w} {c_h}') + + with open(os.path.join(labels_dir, f'{Path(label).stem}.txt'), 'w') as f: + f.writelines(lines) + f.close() + + +if __name__ == '__main__': + convert_xml('datasets/others/UAVimages') \ No newline at end of file