YOLO系列 --- xml2txt脚本

YOLO系列 — xml2txt脚本

YOLO算法的数据格式基本上都是VOC格式的，我在网上看了很多xml2txt的脚本，发现很多还是有一些问题的，并不能直接运行得到正确的txt文件。所以特地写一篇blog~

一、LabelImg会出现的问题

一般来说，VOC格式的图片一般都是用LabelImg作为标注工具进行人工标注的。但是不知道是这个软件还是标注过程，会出现以下几个小bug：

1. 标注生成的xml文件，图片的长宽大小为0。这个问题我一直没找到问题的所在，不过更倾向于图片读取的时候出现异常。
2. 标注框的大小存在越界。这个问题貌似是基于LabelImg版本所导致的，有些版本的LabelImg就会出现标注框的xmin和ymin小于0，xmax和ymax会大于图片的长宽。最终就会导致转换后的txt出现负值或者异常值。

二、xml2txt

话不多数，直接上代码：

import glob
import os
import cv2
import shutil
import numpy as np
import xml.etree.ElementTree as ET


#  转换一个xml文件为txt
def single_xml_to_txt(xml_file, corr_img_path, save_txt_path):
    if os.path.exists(corr_img_path):
        try:
            img = cv2.imread(corr_img_path)
            true_width, true_height = img.shape[1], img.shape[0]
        except:
            # print(img_path)
            pass

    tree = ET.parse(xml_file)
    root = tree.getroot()
    # 保存的txt文件路径
    txt_file = save_txt_path + os.path.basename(xml_file.replace("xml", "txt"))
    with open(txt_file, 'w') as txt_file:
        try:
            picture_width = int(root.find('size')[0].text)
            picture_height = int(root.find('size')[1].text)
            #针对图片尺寸读取异常，进行尺寸修正
            if picture_width == 0 or picture_height == 0:
                picture_width = true_width
                picture_height = true_height
            if picture_width != true_width or picture_height != true_height:
                picture_width = true_width
                picture_height = true_height
            for member in root.findall('object'):
                class_name = member.find("name").text
            	box_x_min = int(member.find("bndbox").find("xmin").text)
            	box_y_min = int(member.find("bndbox").find("ymin").text)
            	box_x_max = int(member.find("bndbox").find("xmax").text)
            	box_y_max = int(member.find("bndbox").find("ymax").text)
				#针对标注框大小异常，进行截断操作，避免转换txt值异常
                if box_x_min < 0:
                    box_x_min = 0
                elif box_y_min < 0:
                    box_y_min = 0
                elif box_x_max > picture_width:
                    box_x_max = picture_width
                elif box_y_max > picture_height:
                    box_y_max = picture_height
                #  类名对应的index
                class_num = class_names.index(class_name)
                # 转成相对位置和宽高
                x_center = (box_x_min + box_x_max) / (2 * picture_width)
                y_center = (box_y_min + box_y_max) / (2 * picture_height)
                width = (box_x_max - box_x_min) / picture_width
                height = (box_y_max - box_y_min) / picture_height
                txt_file.write(str(class_num) + ' ' + str(x_center) + ' ' + str(y_center) + ' ' + str(width) + ' ' + str(height) + '\n')
        except:
            pass

#  转换文件夹下的所有xml文件为txt
def dir_xml_to_txt(xml_path, img_path, save_txt_path):
    for index, xml_file in enumerate(glob.glob(xml_path+ '*.xml')):
        corr_img_path = img_path + os.path.basename(xml_file).split(".")[0] + ".jpg"
        try:
            single_xml_to_txt(xml_file, corr_img_path, save_txt_path)
        except:
            print(xml_file)
        else:
            print("all:{}---now:{}".format(len(glob.glob(xml_path + '*.xml')), index+1))

if __name__ == '__main__':
    class_names = [''] #类名
    xml_path = r"" #xml文件路径
    img_path = r"" #图片路径
    save_txt_path = r"" #生成txt文件保存路径
    dir_xml_to_txt(xml_path, img_path, save_txt_path)