DroneVenicle数据集去除白边并将xml转换成YOLO格式（txt)

最新推荐文章于 2025-08-25 09:00:12 发布

原创最新推荐文章于 2025-08-25 09:00:12 发布 · 1.3k 阅读

38 ·

CC 4.0 BY-SA版权

文章标签：

#python #DroneVehicle

部署运行你感兴趣的模型镜像

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）

一、数据集介绍

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。

二、数据集处理目标

数据集处理的目标有两个：

由于DroneVehicle数据集中的每张图像的顶部、底部、左侧和右侧设置了一个宽度为100像素的白色边框，因此需要生成去除周围白色边框的图像。
由于DroneVehicle数据集的标注采用XML格式进行保存，实验需要TXT格式的标注，因此需要对相关数据进行合理转化

三、去除周围白色边框部分

# Description: This script is used to crop the images in the dataset.
#  我们需要将DroneVehicle数据集裁剪到640*640
import numpy as np
import cv2
import os
from tqdm import tqdm


def create_file(output_dir_vi, output_dir_ir):
    if not os.path.exists(output_dir_vi):
        os.makedirs(output_dir_vi)
    if not os.path.exists(output_dir_ir):
        os.makedirs(output_dir_ir)
    print(f'Created folder:({output_dir_vi}); ({output_dir_ir})')


def update(input_img_path, output_img_path):
    image = cv2.imread(input_img_path)
    cropped = image[100:612, 100:740]  # 裁剪坐标为[y0:y1, x0:x1]
    cv2.imwrite(output_img_path, cropped)


# dataset_dir_vi = '/home/mjy/ultralytics/datasets/OBB/images/val'
# output_dir_vi = '/home/mjy/ultralytics/datasets/OBBCrop/images/val'
# dataset_dir_ir = '/home/mjy/ultralytics/datasets/OBB/image/val'
# output_dir_ir = '/home/mjy/ultralytics/datasets/OBBCrop/image/val'

dataset_dir_vi = 'G:/datasets/DroneVehicle/val/valimg'
output_dir_vi = 'G:/datasets/DroneVehicle/DroneVehicle_crop/val/valimg'
dataset_dir_ir = 'G:/datasets/DroneVehicle/val/valimgr'
output_dir_ir = 'G:/datasets/DroneVehicle/DroneVehicle_crop/val/valimgr'



# 检查文件夹是否存在，如果不存在则创建
create_file(output_dir_vi, output_dir_ir)
# 获得需要转化的图片路径并生成目标路径
image_filenames_vi = [(os.path.join(dataset_dir_vi, x), os.path.join(output_dir_vi, x))
                      for x in os.listdir(dataset_dir_vi)]

image_filenames_ir = [(os.path.join(dataset_dir_ir, x), os.path.join(output_dir_ir, x))
                      for x in os.listdir(dataset_dir_ir)]
# 转化所有图片
print('Start transforming vision images...')

for path in tqdm(image_filenames_vi):
    update(path[0], path[1])


print('Start transforming infrared images...')
for path in tqdm(image_filenames_ir):
    update(path[0], path[1])

dataset_dir_vi：用于存放可见光未处理的图像
output_dir_vi：用于存放可见光处理后的图像
dataset_dir_ir：用于存放红外未处理的图像
output_dir_ir：用于存放红外处理后的图像
可以根据自己的目录进行更改

四、XML数据转TXT数据可视化展示

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

# 图像类别  https://zhuanlan.zhihu.com/p/665126340
classes = [ "car", "truck", "bus", "van", "freight car"]


# 定义相关地址参数
def parse_args():
    parser = parser = argparse.ArgumentParser(description='polygon')
    parser.add_argument('--in_xml_vi_dir', default='G:/datasets/DroneVehicle/DroneVehicle/train/trainlabel', help='可见光 XML 文件地址')
    parser.add_argument('--in_xml_ir_dir', default='G:/datasets/DroneVehicle/DroneVehicle/train/trainlabelr', help='红外光 XML 文件地址')
    parser.add_argument('--out_vi_txt_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/train/trainlabel', help='可见光 TXT 文件地址')
    parser.add_argument('--out_ir_txt_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/train/trainlabelr', help='红外光 TXT 文件地址')
    parser.add_argument('--in_vi_img_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/train/trainimg', help='裁切后的 vi-img 地址')
    parser.add_argument('--in_ir_img_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/train/trainimgr', help='裁切后的 ir-img 地址')


    # parser.add_argument('--in_xml_vi_dir', default='G:/datasets/DroneVehicle/DroneVehicle/val/vallabel', help='可见光 XML 文件地址')
    # parser.add_argument('--in_xml_ir_dir', default='G:/datasets/DroneVehicle/DroneVehicle/val/vallabelr', help='红外光 XML 文件地址')
    # parser.add_argument('--out_vi_txt_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/val/vallabel', help='可见光 TXT 文件地址')
    # parser.add_argument('--out_ir_txt_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/val/vallabelr', help='红外光 TXT 文件地址')
    # parser.add_argument('--in_vi_img_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/val/valimg', help='裁切后的 vi-img 地址')
    # parser.add_argument('--in_ir_img_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/val/valimgr', help='裁切后的 ir-img 地址')

    # parser.add_argument('--in_xml_vi_dir', default='G:/datasets/DroneVehicle/DroneVehicle/test/testlabel', help='可见光 XML 文件地址')
    # parser.add_argument('--in_xml_ir_dir', default='G:/datasets/DroneVehicle/DroneVehicle/test/testlabelr', help='红外光 XML 文件地址')
    # parser.add_argument('--out_vi_txt_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/test/testlabel', help='可见光 TXT 文件地址')
    # parser.add_argument('--out_ir_txt_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/test/testlabelr', help='红外光 TXT 文件地址')
    # parser.add_argument('--in_vi_img_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/test/testimg', help='裁切后的 vi-img 地址')
    # parser.add_argument('--in_ir_img_dir', default='G:/datasets/DroneVehicle/DroneVehicle_crop/test/testimgr', help='裁切后的 ir-img 地址')

    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 == "freight 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' or "feright":  # feright car标签名称在xml文件里好乱呀，有的是是feright car，有的是feright有的是feright_car 统一改成freight car
                name = 'freight 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) - 100
                y1 = int(polygon.find('y1').text) - 100
                x2 = int(polygon.find('x2').text) - 100
                y2 = int(polygon.find('y2').text) - 100
                x3 = int(polygon.find('x3').text) - 100
                y3 = int(polygon.find('y3').text) - 100
                x4 = int(polygon.find('x4').text) - 100
                y4 = int(polygon.find('y4').text) - 100
                # 将获取后的数据填入空列表中
                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) - 100
                y1 = int(ymin) - 100
                x3 = int(xmax) - 100
                y3 = int(ymax) - 100
                # 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)

            # 针对于处理图像后出现负值数据的代码修改
            # 修改思路：在将xml数据转换为txt数据之后，进行一个if条件的判断，如果生成的txt数据有出现负值，则忽略这个数据，否则保存数据
            if 0 <= result[0] <= 1 and 0 <= result[1] <= 1 and 0 <= result[2] <= 1 and 0 <= result[3] <= 1:
                data = str(result[0]) + " " + str(result[1]) + " " + str(result[2]) + " " + str(result[3]) + '\n'
                data = str(result_id) + " " + data
                out_file.write(data)
            else:
                pass


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

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

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

五、可视化展示

import cv2

img = cv2.imread('G:/datasets/DroneVehicle/DroneVehicle_crop_yolo/train/trainimg/00205.jpg')

with open('G:/datasets/DroneVehicle/DroneVehicle_crop_yolo/train/trainlabel/00205.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)

六、转换dota到obb格式

obb方式更注重边框尽可能包围物体而不包含其他背景

hbb方式更注重方向，尽可能水平方向框住物体


# 旋转框数据处理(DroneVehicle数据集)
# 转换Dota数据集格式为YOLO OBB格式
import sys 
from ultralytics.data.converter import convert_dota_to_yolo_obb
convert_dota_to_yolo_obb('/home/mjy/ultralytics/datasets/OBBCrop')