Python 将M3FD数据集格式转为YOLO格式

M3FD从VOC格式转YOLO格式的代码

import os
import xml.etree.ElementTree as ET
import time
from shutil import copyfile
from tqdm import tqdm

"""
本代码这里默认是处理的M3FD数据集
将该数据集由VOC格式转换成YOLO格式
同时将文件重命名存储
同时过滤了自己用不到的类别
使用该代码前，建议稍微读一下，清楚代码中的限制，避免不必要的麻烦
"""
classes = ["People", "Car", "Bus", "Motorcycle", "Lamp", "Truck"]

# 将x1, y1, x2, y2转换成yolov5所需要的x, y, w, h格式
def xyxy2xywh(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[2]) / 2 * dw
    y = (box[1] + box[3]) / 2 * dh
    w = (box[2] - box[0]) * dw
    h = (box[3] - box[1]) * dh
    return (x, y, w, h)         # 返回的都是标准化后的值


def voc2yolo(path):  # 该函数本文件中未用到
    # 可以打印看看该路径是否正确
    print(len(os.listdir(path)))
    # 遍历每一个xml文件
    for file in os.listdir(path):
        # xml文件的完整路径, 注意：因为是路径所以要确保准确，我是直接使用了字符串拼接, 为了保险可以用os.path.join(path, file)
        label_file = path + file
        # 最终要改成的txt格式文件,这里我是放在voc2007/labels/下面
        # 注意: labels文件夹必须存在，没有就先创建，不然会报错
        out_file = open(path.replace('Annotations', 'labels') + file.replace('xml', 'txt'), 'w')
        # print(label_file)

        # 开始解析xml文件
        tree = ET.parse(label_file)
        root = tree.getroot()
        size = root.find('size')            # 图片的shape值
        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
            # 将名称转换为id下标
            cls_id = classes.index(cls)
            # 获取整个bounding box框
            bndbox = obj.find('bndbox')
            # xml给出的是x1, y1, x2, y2
            box = [float(bndbox.find('xmin').text), float(bndbox.find('ymin').text), float(bndbox.find('xmax').text),
                float(bndbox.find('ymax').text)]

            # 将x1, y1, x2, y2转换成yolov5所需要的x_center, y_center, w, h格式
            bbox = xyxy2xywh((w, h), box)
            # 写入目标文件中，格式为 id x y w h
            out_file.write(str(cls_id) + " " + " ".join(str(x) for x in bbox) + '\n')

def save_lab(xml_dir, lab_save, save_name):  # 保存标签文件
    # print(label_file)
    # 开始解析xml文件
    tree = ET.parse(xml_dir)
    root = tree.getroot()
    size = root.find('size')            # 图片的shape值
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    save = False
    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
        # 将名称转换为id下标
        cls_id = classes.index(cls)
        # 注意这里对标签进行了过滤，因为我只需要person类别，所以其他的类别并不需要保存，同时如果xml中不存在person类别也不需要创建文件
        if cls_id == 0:
            out_file = open(os.path.join(lab_save, save_name + ".txt"), 'w')
            save = True
        else:
            continue
        # 获取整个bounding box框
        bndbox = obj.find('bndbox')
        # xml给出的是x1, y1, x2, y2
        box = [float(bndbox.find('xmin').text), float(bndbox.find('ymin').text), float(bndbox.find('xmax').text),
            float(bndbox.find('ymax').text)]

        # 将x1, y1, x2, y2转换成yolov5所需要的x_center, y_center, w, h格式
        bbox = xyxy2xywh((w, h), box)
        # 写入目标文件中，格式为 id x y w h
        out_file.write(str(cls_id) + " " + " ".join(str(x) for x in bbox) + '\n')
    return save


def save_img(img_dir, img_save, save_name):
     copyfile(img_dir, os.path.join(img_save, save_name + ".png"))



if __name__ == '__main__':
    data_path = '/data/jjg/codes/datasets/M3FD_Detection/'
    ann_path = os.path.join(data_path, "Annotation")
    img_path = os.path.join(data_path, "Ir")
    img_save = "/data/jjg/codes/datasets/M3FD_Detection/images/"
    lab_save = "/data/jjg/codes/datasets/M3FD_Detection/labels/"
    xml_list = os.listdir(ann_path)
    for xml_name in tqdm(xml_list):
        xml_dir = os.path.join(ann_path, xml_name)
        img_dir = os.path.join(img_path, xml_name[0:-4] + ".png")
        save_name = str(round(time.time() * 1000)) + "-m3fd"
        save = save_lab(xml_dir, lab_save, save_name)
        if save:
            save_img(img_dir, img_save, save_name)