YOLO格式轉換

import os.path

import xml.etree.ElementTree as ET

class_names = ['car','license-plate'] # your class_names

xmlpath='./xml/' #原xml檔案路徑

txtpath='./yolo/'#轉換後yolo格式(txt檔案)儲存路徑

files = []

for root, dirs, files in os.walk(xmlpath):

    None

if 'desktop.ini' in files :files.remove('desktop.ini')# 移除不相關檔案

number = len(files) #檔案數量

print(number)

i = 0

while i <number:

    name = files[i][0:-4] #移除副檔名(.xml 、 .txt)

    xml_name = name + ".xml"

    txt_name = name + ".txt"

    xml_file_name = xmlpath + xml_name

    txt_file_name = txtpath + txt_name

    print(xml_file_name)

    xml_file = open(xml_file_name)

    tree = ET.parse(xml_file)

    root = tree.getroot()

    filename = root.find('filename').text

    image_name = root.find('filename').text

    w = int(root.find('size').find('width').text)

    h = int(root.find('size').find('height').text)

    f_txt = open(./yolo/txt_file_name, 'w+')

    content = ""

    first = True

    for obj in root.iter('object'):

        name = obj.find('name').text

        class_num = class_names.index(name)

        xmlbox = obj.find('bndbox')

        x1 = int(xmlbox.find('xmin').text)

        x2 = int(xmlbox.find('xmax').text)

        y1 = int(xmlbox.find('ymin').text)

        y2 = int(xmlbox.find('ymax').text)

        if first:

            content += str(class_num) + " " + \

                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \

                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)

            first = False

        else:

            content += "\n" + \

                       str(class_num) + " " + \

                       str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \

                       str((x2 - x1) / w) + " " + str((y2 - y1) / h)

    f_txt.write(content)

    f_txt.close()

    xml_file.close()

    i+=1

print('done')

from lxml import etree

from PIL import Image

import csv

import os

IMG_PATH = "./image/" #圖片資料夾

IMG_FILE = os.listdir(IMG_PATH) #所有圖片檔案名稱

txt_folder = './label/' #yolo格式(.txt)檔案資料夾

labels = ['car','license-plate'] #label名稱

label = ''

def csvread(fn):

    with open(fn, 'r') as csvfile:

        list_arr = []

        reader = csv.reader(csvfile, delimiter=' ')

        for row in reader:

            list_arr.append(row)

    return list_arr

def convert_label(txt_file):

    global label

    for i in range(len(labels)):

        if txt_file == str(i):

            label = labels[i]

            return label

    return label

def extract_coor(txt_file, img_width, img_height):

    x_rect_mid = float(txt_file[1])

    y_rect_mid = float(txt_file[2])

    width_rect = float(txt_file[3])

    height_rect = float(txt_file[4])

    x_min_rect = ((2 * x_rect_mid * img_width) - (width_rect * img_width)) / 2

    x_max_rect = ((2 * x_rect_mid * img_width) + (width_rect * img_width)) / 2

    y_min_rect = ((2 * y_rect_mid * img_height) -

                  (height_rect * img_height)) / 2

    y_max_rect = ((2 * y_rect_mid * img_height) +

                  (height_rect * img_height)) / 2

    return x_min_rect, x_max_rect, y_min_rect, y_max_rect

for line in IMG_FILE:

    root = etree.Element("annotation")

    img_style = IMG_PATH.split('/')[-1]

    img_name = line

    image_info = IMG_PATH + "/" + line

    img_txt_root = txt_folder + "/" + line[:-4]

    txt_path = img_txt_root + ".txt"

    txt_file = csvread(txt_path) #讀取.txt內容

    ######################################

    # 讀取圖片內容

    img_size = Image.open(image_info).size

    img_width = img_size[0]

    img_height = img_size[1]

    img_depth = Image.open(image_info).layers

    ######################################

    folder = etree.Element("folder")

    folder.text = "%s" % (img_style)

    filename = etree.Element("filename")

    filename.text = "%s" % (img_name)

    path = etree.Element("path")

    path.text = "%s" % (IMG_PATH)

    source = etree.Element("source")

    ##################source - element##################

    source_database = etree.SubElement(source, "database")

    source_database.text = "Unknown"

    ####################################################

    size = etree.Element("size")

    ####################size - element##################

    image_width = etree.SubElement(size, "width")

    image_width.text = "%d" % (img_width)

    image_height = etree.SubElement(size, "height")

    image_height.text = "%d" % (img_height)

    image_depth = etree.SubElement(size, "depth")

    image_depth.text = "%d" % (img_depth)

    ####################################################

    segmented = etree.Element("segmented")

    segmented.text = "0"

    root.append(folder)

    root.append(filename)

    root.append(path)

    root.append(source)

    root.append(size)

    root.append(segmented)

    for ii in range(len(txt_file)):

        label = convert_label(txt_file[ii][0])

        x_min_rect, x_max_rect, y_min_rect, y_max_rect = extract_coor(

            txt_file[ii], img_width, img_height)

        object = etree.Element("object")

        ####################object - element##################

        name = etree.SubElement(object, "name")

        name.text = "%s" % (label)

        pose = etree.SubElement(object, "pose")

        pose.text = "Unspecified"

        truncated = etree.SubElement(object, "truncated")

        truncated.text = "0"

        difficult = etree.SubElement(object, "difficult")

        difficult.text = "0"

        bndbox = etree.SubElement(object, "bndbox")

        #####sub_sub########

        xmin = etree.SubElement(bndbox, "xmin")

        xmin.text = "%d" % (x_min_rect)

        ymin = etree.SubElement(bndbox, "ymin")

        ymin.text = "%d" % (y_min_rect)

        xmax = etree.SubElement(bndbox, "xmax")

        xmax.text = "%d" % (x_max_rect)

        ymax = etree.SubElement(bndbox, "ymax")

        ymax.text = "%d" % (y_max_rect)

        #####sub_sub########

        root.append(object)

        ####################################################

    file_output = etree.tostring(root, pretty_print=True, encoding='UTF-8')

    # print(file_output.decode('utf-8'))

    ff = open(f'./xml/{img_name[:-4]}.xml' , 'w', encoding="utf-8")

    ff.write(file_output.decode('utf-8'))