近期用到无人机UAVDT,需要在yolo上进行训练,下载后分为多目标和单目标两种,本次使用多目标数据进行训练,即对UAV-benchmark-MOTD_v1.0中的xx_gt_whole.txt真值文件进行修改。

查看了几篇修改过程:将UAVDT数据集格式转成VOC、YOLO格式(最终版)_uavdt数据集没标签-CSDN博客

将UAVDT数据集转化成YOLO样式-CSDN博客

都要先转换成COCO再到VOC最后再到yolo,不太适用于我,于是根据他们的教程进行修改,以下是我的转换过程:

1.处理数据集中的图像和相应的ground truth标签,将xx_gt_whole.txt按照图像文件的顺序保存到相应的目录中。生成目录如下:

def process_dataset(data_path):
    data_num = 0
    for name in os.listdir(data_path):

       # print('name:', name)
        data_num += 1

        img_path = os.path.join(data_path, name)
        gt_path = os.path.join(data_path.replace('UAV-benchmark-M', 'UAV-benchmark-MOTD_v1.0/GT'), name + '_gt_whole.txt')
        save_gt_path = os.path.join(data_path, name, 'gt')

        photo_count = 0
        for file_name in os.listdir(img_path):
            _, extension = os.path.splitext(file_name)
            if extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
                photo_count += 1
        print("Total photos:", photo_count)

        if not os.path.exists(save_gt_path):
            os.makedirs(save_gt_path)

        with open(gt_path, 'r') as file:
            lines = file.readlines()
            for i in range(1, photo_count + 1):
                file_path1 = os.path.join(save_gt_path, str(i) + '.txt')
                with open(file_path1, 'w') as target_file:
                    for line in lines:
                        data = line.split(',')
                        if data[0] == str(i):
                            target_file.write(line)
   # print('img_num={}\nend!!!'.format(data_num))

2.将所有图片和标注分别放在/img和/gt的文件夹中,并修改为一一对应的名字。 

def img_and_gt(data_path, save_path):
    os.makedirs(save_path, exist_ok=True)
    # 目标图片文件夹路径
    save_img_folder = os.path.join(save_path, 'img')
    os.makedirs(save_img_folder, exist_ok=True)
    # 目标标签文件夹路径
    save_gt_folder = os.path.join(save_path, 'gt')
    os.makedirs(save_gt_folder, exist_ok=True)

    # 遍历指定文件夹中的所有文件夹
    for folder_name in tqdm(os.listdir(data_path)):
        # 构建当前文件夹中的图片文件夹路径
        img_folder_path = os.path.join(data_path, folder_name)
        # 用于计数的变量
        img_count = 1
        gt_count = 1
        # 遍历图片文件夹中的所有文件
        for filename in os.listdir(img_folder_path):
            # 获取文件的扩展名
            _, extension = os.path.splitext(filename)
            if extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
                # 获取文件的完整路径
                file_path = os.path.join(img_folder_path, filename)
                # 生成新的文件名
                new_filename = f"{folder_name}_{img_count}{extension}"
                # 构建目标文件的完整路径
                destination_file = os.path.join(save_img_folder, new_filename)
                # 移动文件到目标文件夹
                # print(file_path,'/n',destination_file)
                shutil.copy(file_path, destination_file)
                # 更新计数器
                # print('num:',img_count)
                img_count += 1

        # 构建当前文件夹中的标签文件夹路径
        gt_folder_path = os.path.join(data_path, folder_name, "gt")
        # 如果标签文件夹存在
        if os.path.exists(gt_folder_path):
            dir = os.listdir(gt_folder_path)
            dir.sort(key=lambda x: int(x[:-4]))
            # 遍历标签文件夹中的所有文件
            for filename in dir:
                # 获取文件的完整路径
                file_path = os.path.join(gt_folder_path, filename)
                # 生成新的文件名
                new_filename = f"{folder_name}_{gt_count}.txt"
                # 构建目标文件的完整路径
                destination_file = os.path.join(save_gt_folder, new_filename)
                # 移动文件到目标文件夹
                shutil.copy(file_path, destination_file)
                # 更新计数器
                gt_count += 1

3.将所有图片分为训练集和测试集。

def split_dataset(source_dir, train_dir, test_dir, split_ratio=0.8):
    # 创建训练集和测试集文件夹
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 获取所有图片文件的列表
    image_files = [f for f in os.listdir(source_dir) if f.endswith('.jpg') or f.endswith('.png')]

    # 随机化文件列表
    random.shuffle(image_files)

    # 计算训练集和测试集的数量
    num_train = int(len(image_files) * split_ratio)
    num_test = len(image_files) - num_train

    # 将文件移动到相应的文件夹中
    for i, image_file in tqdm(enumerate(image_files)):
        source_path = os.path.join(source_dir, image_file)
        if i < num_train:
            target_path = os.path.join(train_dir, image_file)
        else:
            target_path = os.path.join(test_dir, image_file)
        shutil.copy(source_path, target_path)

4.为训练集分配对应的标签文件。

def copy_txt_files(image_folder, txt_folder, output_folder):
    # 创建输出文件夹
    os.makedirs(output_folder, exist_ok=True)
    # 获取图片文件夹中所有图片文件的名称列表
    image_files = [f for f in os.listdir(image_folder) if f.endswith('.jpg') or f.endswith('.png')]
    # 遍历图片文件夹中的图片文件
    for image_file in tqdm(image_files) :
        # 构建对应的txt文件路径
        txt_file = os.path.splitext(image_file)[0] + '.txt'
        txt_path = os.path.join(txt_folder, txt_file)
        # 检查对应的txt文件是否存在
        if os.path.exists(txt_path):
            # 复制txt文件到输出文件夹中
            shutil.copy(txt_path, output_folder)

5.将标签文件格式转为yolo格式。

def convert_uavdt_to_yolo(img_folder, uavdt_labels_folder, yolo_labels_folder):
    for uavdt_file in tqdm(os.listdir(uavdt_labels_folder)):
        if uavdt_file.endswith('.txt'):
            uavdt_labels_file = os.path.join(uavdt_labels_folder, uavdt_file)
            os.makedirs(yolo_labels_folder, exist_ok=True)
            yolo_labels_file = os.path.join(yolo_labels_folder, uavdt_file)
            name, _ = os.path.splitext(uavdt_file)
            img_file = os.path.join(img_folder, f"{name}.jpg")
            with Image.open(img_file) as img:
                img_width, img_height = img.size
            with open(uavdt_labels_file, 'r') as f_in, open(yolo_labels_file, 'w') as f_out:
                for line in f_in:
                    data = line.strip().split(',')
                    frame_index, target_id, bbox_left, bbox_top, bbox_width, bbox_height, out_of_view, occlusion, object_category = data

                    x_center = (float(bbox_left) + float(bbox_width) / 2) / img_width
                    y_center = (float(bbox_top) + float(bbox_height) / 2) / img_height
                    width = float(bbox_width) / img_width
                    height = float(bbox_height) / img_height

                    yolo_format = f"{object_category} {x_center} {y_center} {width} {height}"
                    f_out.write(yolo_format + '\n')

6.查看结果是否正确。

def visualize_yolo_labels(image_path, yolo_labels_path):
    # 读取图像
    image = cv2.imread(image_path)
    print(image.shape)
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 读取YOLO格式的标签文件
    with open(yolo_labels_path, 'r') as f:
        for line in f:
            data = line.strip().split(' ')
            object_category, x_center, y_center, width, height = data

            # 计算边界框坐标
            x_center, y_center, width, height = map(float, [x_center, y_center, width, height])
            img_h, img_w, _ = image.shape
            x_center *= img_w
            y_center *= img_h
            width *= img_w
            height *= img_h
            x1, y1, x2, y2 = int(x_center - width/2), int(y_center - height/2), int(x_center + width/2), int(y_center + height/2)

            # 绘制边界框
            cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
            cv2.putText(image, object_category, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

总共有三类:1:'car',2:'truck',3:'bus' 

使用时:

# 调用函数
data_path = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M'
print('start step 1')
process_dataset(data_path)

save_path = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo'
print('start step 2')
img_and_gt(data_path, save_path)

img_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\img'
train_img_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\train_img'
test_img_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\test_img'
print('start step 3')
split_dataset(img_dir, train_img_dir, test_img_dir)

txt_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\gt'
train_gt_uav_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\train_gt_uav'
print('start step 4')
copy_txt_files(train_img_dir, txt_dir, train_gt_uav_dir)

train_gt_yolo_dir = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\train_gt_yolo'
print('start step 5')
convert_uavdt_to_yolo(train_img_dir, train_gt_uav_dir, train_gt_yolo_dir)

image_path = r'data/UAV-benchmark-M-yolo/train_img/M0201_181.jpg'
yoly_labels_path = r'E:\无人机数据集\UAVDT\data\UAV-benchmark-M-yolo\train_gt_yolo\M0201_181.txt'
visualize_yolo_labels(image_path, yoly_labels_path)

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐