YOLO系列笔记（七）——xml文件转换成txt文件

YOLO系列笔记（七）——xml文件转换成txt文件

该笔记将介绍将上一篇笔记中用labelimg标注出的文.xml文件转换成.txt文件的代码，其中大部分来自这位作者的这篇博客分享，笔者仅对改代码进行些许修改和介绍分析，在此十分感谢其作者的帮助和分享！

代码

import os
import glob
import argparse
import random
import xml.etree.ElementTree as ET
from PIL import Image
from tqdm import tqdm

def get_all_classes(xml_path):
    xml_fns = glob.glob(os.path.join(xml_path, '*.xml'))
    class_names = []
    for xml_fn in xml_fns:
        tree = ET.parse(xml_fn)
        root = tree.getroot()
        for obj in root.iter('object'):
            cls = obj.find('name').text
            class_names.append(cls)
    return sorted(list(set(class_names)))

def convert_annotation(img_path, xml_path, class_names, out_path):
    output = []
    im_fns = glob.glob(os.path.join(img_path, '*.jpg'))
    for im_fn in tqdm(im_fns):
        if os.path.getsize(im_fn) == 0:
            continue
        xml_fn = os.path.join(xml_path, os.path.splitext(os.path.basename(im_fn))[0] + '.xml')
        if not os.path.exists(xml_fn):
            continue
        img = Image.open(im_fn)
        height, width = img.height, img.width
        tree = ET.parse(xml_fn)
        root = tree.getroot()
        anno = []
        xml_height = int(root.find('size').find('height').text)
        xml_width = int(root.find('size').find('width').text)
        if height != xml_height or width != xml_width:
            print((height, width), (xml_height, xml_width), im_fn)
            continue
        for obj in root.iter('object'):
            cls = obj.find('name').text
            cls_id = class_names.index(cls)
            xmlbox = obj.find('bndbox')
            xmin = int(xmlbox.find('xmin').text)
            ymin = int(xmlbox.find('ymin').text)
            xmax = int(xmlbox.find('xmax').text)
            ymax = int(xmlbox.find('ymax').text)
            cx = (xmax + xmin) / 2.0 / width
            cy = (ymax + ymin) / 2.0 / height
            bw = (xmax - xmin) * 1.0 / width
            bh = (ymax - ymin) * 1.0 / height
            anno.append('{} {} {} {} {}'.format(cls_id, cx, cy, bw, bh))
        
        if len(anno) > 0:
            output.append(im_fn)
            txt_out_path = 'data/labels/'
            # 构建新的 .txt 文件的完整路径
            txt_filename = os.path.basename(im_fn).replace('.jpg', '.txt')  # 从图像文件名生成文本文件名
            txt_file_path = os.path.join(txt_out_path, txt_filename)  # 组合目录路径和文件名

            # 使用新路径保存文件
            with open(txt_file_path, 'w') as f:
                f.write('\n'.join(anno))

    random.shuffle(output)
    train_num = int(len(output) * 0.9)
    with open(os.path.join(out_path, 'train.txt'), 'w') as f:
        f.write('\n'.join(output[:train_num]))
    with open(os.path.join(out_path, 'val.txt'), 'w') as f:
        f.write('\n'.join(output[train_num:]))

def parse_args():
    parser = argparse.ArgumentParser('generate annotation')
    parser.add_argument('--img_path', type=str, help='input image directory')
    parser.add_argument('--xml_path', type=str, help='input xml directory')
    parser.add_argument('--out_path', type=str, help='output directory')
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    class_names = get_all_classes(args.xml_path)
    print(class_names)
    convert_annotation(args.img_path, args.xml_path, class_names, args.out_path)

分析

这段代码实现了一个用于处理图像及其对应XML注释的脚本，整理了所有图像文件中出现的类别标签，将.xml文件中关于框选类别和尺寸的信息整理成了.txt文件，并将所含有图像文件分成了用于训练和测试用途的两部分。

导入的库

• os：用于处理文件和目录路径。
• glob：用于查找匹配特定模式的文件路径名。
• argparse：用于解析命令行参数。
• random：用于生成随机数，这里主要用于打乱数据。
• xml.etree.ElementTree（ET）：用于解析和创建XML数据。
• PIL.Image：用于图像处理。
• tqdm：用于在循环操作中显示进度条。

函数说明

1. get_all_classes(xml_path)

   • 功能：从给定目录下的XML文件中提取所有类别名称。
   • 实现：使用glob.glob查找所有XML文件，然后使用ET包解析每个文件，找到每个对象的类别名称。最终返回所有独特类别的排序列表。

2. convert_annotation(img_path, xml_path, class_names, out_path)

   • 功能：转换xml文件为txt格式，并生成训练和验证集。
   • 实现：找到img_path下所有图像文件，对它们进行遍历，遍历的时候通过tqdm包提供了用户友好的进度显示。在每层遍历中，首先对图像文件大小为零的情况进行了检查，如果为零，则跳过这样的文件。其次对于每张图像检查对应的XML文件是否存在，并验证图像尺寸与XML中记录的尺寸是否一致。如果尺寸一致，则检查该图像中的所有标签类别，和计算对应对象（也就是框选对象）的中心坐标和宽高比例（相对于图像尺寸），生成可用的标注信息，最后将所有标注信息写入与图像同名的.txt文件中。完成后，随机打乱所有处理过的图像路径，按9:1的比例分为训练集和验证集，并保存到相应的文本文件中。注意，在这段代码中，将所有存有标注信息的.txt文件都存储到了相对路径‘data/labels’中，如果想要自行更改，可以在该函数的txt_out_path中更改，或者在下面的parse_args()函数中增加该变量，在执行整个Python文件的时候进行输入。

3. parse_args()

   • 功能：解析命令行输入的参数。
   • 实现：设置了三个命令行参数- -img_path，- -xml_path，和- -out_path，分别用于指定图像目录、XML目录和输出目录。

4. 主执行流程

   • 在脚本的主体部分，首先解析命令行参数，然后使用这些参数调用get_all_classes函数和convert_annotation函数。首先获取所有类别名称，并打印出来，然后开始转换注释并生成训练和验证集文件。

执行命令

以下对于改代码的执行命令：

python generate_txt.py --img_path data/images --xml_path data/Annotations --out_path data/label

其中位于首列的‘python’要根据自己系统上的Python版本进行更改，如果你安装的python版本较高，需要把这一列换成python3。另外img_path，xml_path和out_path也要根据自己的文件安排进行更改，这里重申一下out_path是用来存放训练集和测试集的路径，存放标签的.txt文件集路径在代码中被定义为了‘data/labels’，该路径可以在代码中被更改，也可以通过命令行参数传递进来，那样代码一需要一定的修改。博主就是在路径这个地方栽了好大的跟头，希望大家多多注意。

最后，看到这里如果觉得该笔记对您有用的话，可以点个小小的赞，或者点赞收藏关注一键三连ヾ(◍’౪`◍) ~ 谢谢！！