Convert VOC annotations to the YOLO format

voc_to_yolo.py

from typing import List
import xml.etree.ElementTree as ET
import os
import sys

from pathlib import Path


def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x, w, y, h = x*dw,  w*dw,  y*dh,  h*dh
    return (x, y, w, h)


def convert_annotation(file: Path, output: Path, classes: List[str]):
    out_file = open(f"{output}/{file.stem}.txt", "w")
    tree = ET.parse(file)
    root = tree.getroot()
    size = root.find("size")
    w = int(size.find("width").text)
    h = int(size.find("height").text)
    for obj in root.iter("object"):
        difficult = obj.find("difficult").text
        cls = obj.find("name").text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find("bndbox")
        b = (float(xmlbox.find("xmin").text), float(xmlbox.find("xmax").text), float(
            xmlbox.find("ymin").text), float(xmlbox.find("ymax").text))
        bb = convert((w, h), b)
        result = f'{cls_id} {" ".join([str(a) for a in bb])}\n'
        out_file.write(result)


def convert_annotations(path: str, output: str, classes: List[str]):
    for filename in os.listdir(path):
        if not filename.endswith(".xml"):
            continue
        file = Path(os.path.join(path, filename))
        if not os.path.exists(output):
            os.makedirs(output)
        convert_annotation(file, Path(output), classes)


if __name__ == "__main__":
    source, output, classes = sys.argv[1], sys.argv[2], sys.argv[3].split(",")
    convert_annotations(source, output, classes)