Conjunto de datos Rope3D en formato YOLO para detección de objetos 2D

1. Descarga del conjunto de datos

Dirección del sitio web oficial: https://thudair.baai.ac.cn/index

2. Introducción al formato de archivo de anotación del conjunto de datos Rope3D

标注数据：
car -1 -1 -10 166.00 111.00 181.00 127.00 -1 -1 -1 -1000 -1000 -1000 -10
参数定义：
Type X X X left_top_x left_top_y right_bottom_x right_bottom_y X X X X X X X

Solo es necesario utilizar los archivos de anotación de YOLO,izquierda_arriba_x, izquierda_arriba_y, derecha_abajo_x, derecha_abajo_y, otros parámetros se pueden ignorar directamente.

3. Código de procesamiento

from PIL import Image
import os


# 文件格式为txt
# 1 -1 -1 -10 166.00 111.00 181.00 127.00 -1 -1 -1 -1000 -1000 -1000 -10
# Label:
# Type X X X left_top_x left_top_y right_bottom_x right_bottom_y X X X X X X X
# 需要读取每行数据，并坐标信息进行转化

def convert(size, box):

    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]

    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]

    # print(x, y, w, h)
    return (x, y, w, h)


# 获取文件夹下所有图片的高和宽信息，并通过字典返回
def getImageWHInfo():
    # 保存数据的字典
    image_info_dict = {
    
    }
    # 获取文件夹中所有文件的名称
    filenames = os.listdir("E:\project\Rope3D dataset\original/train_images")
    for filename in filenames:
        # 打开图片文件
        image = Image.open("E:\project\Rope3D dataset\original/train_images/" + filename)

        # 获取图片的宽和高
        width, height = image.size
        # 保存图片信息
        image_info_dict[filename] = [width, height]

        # 关闭图片
        image.close()

    return image_info_dict


if __name__ == '__main__':
    # 获取文件夹下所有的图片的width，high信息
    image_info_dict = getImageWHInfo()
    save_txt_files_path = 'E:\project\Rope3D dataset\original/train_labels_yolo'
    success_num = 0
    total = len(image_info_dict)
    # Rope3D数据集实际分类不只有8类，我进行了筛选，通过下面的分类过滤代码。如果不需要，可以修改为官方类别，并注释掉下面的过滤代码。
    category_label = {
    
    'car': 0, 'truck': 1, 'van': 2, 'bus': 3, 'pedestrian': 4, 'cyclist': 5, 'tricyclist': 6,
                     'motorcyclist': 7}
    for fileName, WHList in image_info_dict.items():
        print(fileName, WHList)

        image_width = WHList[0]
        image_height = WHList[1]

        fileName = fileName.split('.')[0] + '.txt'

        out_txt_path = os.path.join(save_txt_files_path, fileName)
        out_txt_f = open(out_txt_path, 'w')

        # 打开文件
        with open('E:\project\Rope3D dataset\original/training\label_2/' + fileName, 'r') as f:
            # 逐行读取文件内容
            for line in f:
                # 去除行尾的空白字符（包括空格和换行符）
                line = line.strip()
                # 如果行不为空，则输出该行
                if line:
                    line_list = line.split()
                    category = line_list[0]
                    # 分类过滤代码，可以过滤掉不想要分类。
                    if category in ['barrow', 'unknowns_movable', 'trafficcone', 'unknown_unmovable']:
                        continue
                    categoryLabel = category_label[category]
                    xmin = left_top_x = line_list[4]
                    ymin = left_top_y = line_list[5]
                    xmax = right_bottom_x = line_list[6]
                    ymax = right_bottom_y = line_list[7]
                    b = (float(xmin), float(xmax), float(ymin), float(ymax))
                    bb = convert((image_width, image_height), b)
                    out_txt_f.write(str(categoryLabel) + " " + " ".join([str(a) for a in bb]) + '\n')
            success_num = success_num + 1
            print("写入成功:" + str(success_num) + '/' + str(total))

4. Visualización de resultados