xmlファイルからyolo用のtxtに変換する方法

はじめに

至らない点なども多いと思いますが、ご了承いただければと思います。
ここでは、Vottなどでアノテーションを行った際に生成されるアノテーションファイル(.xml)をyoloで使用するテキストファイル(.txt)に置き換えるコードを載せたいと思います。

コードの内容

import xml.etree.ElementTree as ET
import os
from PIL import Image

classes = ['0'] # クラス名のリスト

def convert_annotation(xml_path, txt_path):
    tree = ET.parse(xml_path)
    root = tree.getroot()

    with open(txt_path, 'w') as f:
        for obj in root.iter('object'):
            cls = obj.find('name').text
            if cls not in classes:
                continue
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('xmax').text), float(xmlbox.find('ymax').text))
            bb = convert((width, height), b)
            f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

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

if __name__ == '__main__':
    img_dir = '/content/drive/MyDrive/ナンバープレート-PascalVOC-export/JPEGImages/' # 画像フォルダのパス
    xml_dir = '/content/drive/MyDrive/ナンバープレート-PascalVOC-export/Annotations/' # XMLフォルダのパス
    txt_dir = '/content/drive/MyDrive/ナンバープレート-PascalVOC-export/txt' # TXTフォルダのパス

    if not os.path.exists(txt_dir):
        os.makedirs(txt_dir)

    xml_files = os.listdir(xml_dir)
    for xml_file in xml_files:
        if xml_file.endswith('.xml'):
            img_name = os.path.splitext(xml_file)[0] + '.jpg'
            img_path = os.path.join(img_dir, img_name)
            xml_path = os.path.join(xml_dir, xml_file)
            txt_path = os.path.join(txt_dir, os.path.splitext(xml_file)[0] + '.txt')

            img = Image.open(img_path)
            width, height = img.size

            convert_annotation(xml_path, txt_path)