DroneVehicle数据集标签转换(.xml→.txtyolo水平框)

目录

1.数据集下载:

2.数据集介绍：

3.数据集标签转换

4.数据集标签可视化

---

1.数据集下载:

DroneVenicle数据集是由天津大学收集、标注的大型无人机航拍车辆数据集。
DroneVenicle训练集下载地址：https://pan.baidu.com/s/1ptZCJ1mKYqFnMnsgqEyoGg（密码：ngar）
DroneVenicle验证集下载地址：https://pan.baidu.com/s/1e6e9mESZecpME4IEdU8t3Q（密码：jnj6）
DroneVenicle测试集下载地址：https://pan.baidu.com/s/1JlXO4jEUQgkR1Vco1hfKhg（密码：tqwc）

2.数据集介绍：

DroneVehicle数据集由无人机采集的56,878张图像组成，其中一半是RGB图像，其余为红外图像。

作者为这五个类别制作了带有定向边界框的丰富注释。其中，汽车（car）有389,779个RGB图像注释，428,086个红外图像注释，卡车（truck）有22,123个RGB图像注释，25,960个红外图像注释，公共汽车（bus）有15,333个RGB图像注释，红外图像有16,590个注释，面包车（van）有11,935个RGB图像注释，红外图像有12,708个注释，货车（freight car）有13,400个RGB图像注释，和17,173个红外图像注释。

在DroneVehicle中，为了在图像边界处标注对象，作者在每张图像的顶部、底部、左侧和右侧设置了一个宽度为100像素的白色边框，因此下载的图像比例为840 x 712，去除周围的白色边框并将图像比例更改为640 x 512。

3.数据集标签转换

1、 Dronevehicle 数据集中有旋转框和水平框，转换时要进行分类获取坐标
2、其中有个类别为feright car，我把官方的feright car全部改写为feright_car,所以要记得更改数据集配置文件里类别名字。

import os
import xml.etree.ElementTree as ET
import math
import cv2 as cv
import argparse
from tqdm import tqdm

# 图像类别
classes = ["feright_car", "car", "truck", "bus", "van"]


# 定义相关地址参数
def parse_args():
    parser = parser = argparse.ArgumentParser(description='polygon')
    parser.add_argument('--in_train_xml_vi_dir', default='./datasets/DroneVehicle/train-original/train/trainlabel', help='train的XML 文件地址')
    parser.add_argument('--in_val_xml_vi_dir', default='./datasets/DroneVehicle/val-original/val/vallabelr', help='val的XML 文件地址')
    parser.add_argument('--out_train_txt_vi_dir', default='./datasets/DroneVehicle/labels/train', help='train TXT 输出文件地址')
    parser.add_argument('--out_val_txt_vi_dir', default='./datasets/DroneVehicle/labels/val', help='val TXT 输出文件地址')
    parser.add_argument('--in_train_img_vi_dir', default='./datasets/DroneVehicle/train-original/train/trainimg', help='train 可见光图片地址')
    parser.add_argument('--in_val_img_vi_dir', default='./datasets/DroneVehicle/val-original/val/valimg', help='val 可见光图片地址')

    args = parser.parse_args()
    return args


# 根据 xml 文件中的 name 选择生成的 txt 文件中的 id
def select_id(name):
    if name == "car":
        id = 0
    elif name == "truck":
        id = 1
    elif name == "bus":
        id = 2
    elif name == "van":
        id = 3
    elif name == "feright_car":
        id = 4
    return id


# YOLO 数据处理
def data_transform(height, width, xmin, ymin, xmax, ymax):
    # 中心点坐标 x_c,y_c
    x_c = (xmin + xmax) / 2
    y_c = (ymin + ymax) / 2

    # 中心横坐标与图像宽度比值 x_，中心纵坐标与图像高度比值 y_，bbox 宽度与图像宽度比值 w_，bbox 高度与图像高度比值 h_
    x_ = x_c / width
    y_ = y_c / height
    w_ = (xmax - xmin) / width
    h_ = (ymax - ymin) / height

    return x_, y_, w_, h_


# xml 文件转 txt 文件
def xml2txt(in_xml_dir, xml_name, out_txt_dir, in_img_dir):
    txt_name = xml_name[:-4] + '.txt'   # 获取生成的 txt 文件名
    txt_path = out_txt_dir  # 获取生成的 txt 文件保存地址

    # 判断保存 txt 文件的文件夹是否存在，如果不存在则创建相应文件夹
    if not os.path.exists(txt_path):
        os.makedirs(txt_path)
    txt_file = os.path.join(txt_path, txt_name)     # 获取 txt 文件地址（保存地址 + 保存名字）

    img_name = xml_name[:-4] + '.jpg'   # 获取图像名字，确保生成的 txt 文件名与图像文件名一致
    img_path = os.path.join(in_img_dir, img_name)   # 获取图像地址
    img = cv.imread(img_path)   # 读取图像信息
    height, width, _ = img.shape    # 获取图像高度（height），宽度（width），通道数（_）

    xml_file = os.path.join(in_xml_dir, xml_name)   # 获取 xml 文件地址
    tree = ET.parse(os.path.join(xml_file))     # 使用 ET.parse 方法解析 xml 文件
    root = tree.getroot()   # 使用 getroot 方法获取根目录

    # 生成对应的 txt 文件
    with open(txt_file, "w+", encoding='UTF-8') as out_file:
        for obj in root.findall('object'):
            # 修改部分标注文件中标注不全的 name 文件
            name = obj.find('name').text
            if name == 'feright car':
                name = 'feright_car'
            else:
                name = name

            # 从 xml 文件中提取相关数据信息,并进行删除白边数据操作（白边宽度 100 像素）
            if obj.find('polygon'):
                # 创建空列表用于存放需要处理的数据
                xmin, xmax, ymin, ymax = [], [], [], []
                polygon = obj.find('polygon')
                # 使用 .find() 方法获取对应 xml 文件中键的键值
                x1 = int(polygon.find('x1').text)
                y1 = int(polygon.find('y1').text)
                x2 = int(polygon.find('x2').text)
                y2 = int(polygon.find('y2').text)
                x3 = int(polygon.find('x3').text)
                y3 = int(polygon.find('y3').text)
                x4 = int(polygon.find('x4').text)
                y4 = int(polygon.find('y4').text)
                # 将获取后的数据填入空列表中
                for i in [x1, x2, x3, x4]:
                    xmin.append(i)
                    xmax.append(i)
                for j in [y1, y2, y3, y4]:
                    ymin.append(j)
                    ymax.append(j)
                # 使用 min()、max() 方法获取最大值，最小值
                xmin = min(xmin)
                xmax = max(xmax)
                ymin = min(ymin)
                ymax = max(ymax)
                # yolo 格式转换
                result = data_transform(height, width, xmin, ymin, xmax, ymax)
                # id 选择
                result_id = select_id(name)

            elif obj.find('bndbox'):
                bndbox = obj.find('bndbox')
                # 使用 .find() 方法获取对应 xml 文件中键的键值
                xmin = bndbox.find('xmin').text
                ymin = bndbox.find('ymin').text
                xmax = bndbox.find('xmax').text
                ymax = bndbox.find('ymax').text
                x1 = int(xmin)
                y1 = int(ymin)
                x3 = int(xmax)
                y3 = int(ymax)
                # yolo 格式转换
                result = data_transform(height, width, x1, y1, x3, y3)
                # id 选择
                result_id = select_id(name)

            # 创建 txt 文件中的数据
            data = str(result[0]) + " " + str(result[1]) + " " + str(result[2]) + " " + str(result[3]) + '\n'
            data = str(result_id) + " " + data
            out_file.write(data)


if __name__ == "__main__":
    args = parse_args()     # 获取命令参数
    xml_vi_path = args.in_train_xml_vi_dir    # 获取可见光 train xml 文件地址
    xmlFiles_vi = os.listdir(xml_vi_path)   # 生成可见光 train xml 文件名列表
    xml_ir_path = args.in_val_xml_vi_dir    # 获取可见光 val xml 文件地址
    xmlFiles_ir = os.listdir(xml_ir_path)   # 生成可见光 val xml 文件名列表

    print('Start transforming vision labels...')
    for i in tqdm(range(0, len(xmlFiles_vi))):
        xml2txt(args.in_train_xml_vi_dir, xmlFiles_vi[i], args.out_train_txt_vi_dir, args.in_train_img_vi_dir)
    print('Finish.')

    print('Start transforming infrared labels...')
    for i in tqdm(range(0, len(xmlFiles_ir))):
        xml2txt(args.in_val_xml_vi_dir, xmlFiles_ir[i], args.out_val_txt_vi_dir, args.in_val_img_vi_dir)
    print('Finish.')

4.数据集标签可视化

import cv2

img = cv2.imread('./datasets/DroneVehicle/images/train/00001.jpg')

with open('./datasets/DroneVehicle/labels/train/00001.txt', 'r') as f:
    for line in f.readlines():
        # temp = f.read()
        temp = line.split()
        x_, y_, w_, h_ = eval(temp[1]), eval(temp[2]), eval(temp[3]), eval(temp[4])

        x1, y1, x2, y2 = int((x_ - w_/2) * img.shape[1]), int((y_ - h_/2) * img.shape[0]), \
                         int((x_ + w_/2) * img.shape[1]), int((y_ + h_/2) * img.shape[0])
        cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0))

cv2.imshow('windows', img)
cv2.waitKey(0)