CUB200-2011数据集转换为YOLO格式

 

数据集的文件结构如下：

CUB_200_2011/
|-- images/
|-- attributes/
|-- bounding_boxes.txt
|-- image_class_labels.txt
|-- images.txt
|-- train_test_split.txt
 

• images/: 包含所有图像。
• bounding_boxes.txt: 每行格式为 <image_id> <x> <y> <width> <height>，表示图像的边界框。
• image_class_labels.txt: 每行格式为 <image_id> <class_id>，表示图像的类别。
• images.txt: 每行格式为 <image_id> <image_path>，表示图像的文件路径。
• train_test_split.txt: 每行格式为 <image_id> <is_training_image>，1表示训练集，0表示测试集。

YOLO格式要求每个图像有一个对应的文本文件，包含以下格式的行：

<class_id> <x_center> <y_center> <width> <height>
 

转换脚本：

import os
import pandas as pd

# 路径设置
dataset_path = 'CUB_200_2011'
images_path = os.path.join(dataset_path, 'images')
output_path = 'yolo_labels'

# 创建输出目录
os.makedirs(output_path, exist_ok=True)

# 读取文件
bounding_boxes = pd.read_csv(os.path.join(dataset_path, 'bounding_boxes.txt'), sep=' ', header=None, names=['image_id', 'x', 'y', 'width', 'height'])
image_class_labels = pd.read_csv(os.path.join(dataset_path, 'image_class_labels.txt'), sep=' ', header=None, names=['image_id', 'class_id'])
images = pd.read_csv(os.path.join(dataset_path, 'images.txt'), sep=' ', header=None, names=['image_id', 'image_path'])
train_test_split = pd.read_csv(os.path.join(dataset_path, 'train_test_split.txt'), sep=' ', header=None, names=['image_id', 'is_training'])

# 将类别ID从1开始调整为从0开始
image_class_labels['class_id'] -= 1

# 合并数据
data = pd.merge(bounding_boxes, image_class_labels, on='image_id')
data = pd.merge(data, images, on='image_id')

# YOLO标签生成
for _, row in data.iterrows():
    image_id = row['image_id']
    class_id = row['class_id']
    image_path = row['image_path']
    x, y, width, height = row['x'], row['y'], row['width'], row['height']

    # 读取图像尺寸
    full_image_path = os.path.join(images_path, image_path)
    img = cv2.imread(full_image_path)
    img_h, img_w = img.shape[:2]

    # 转换为YOLO格式
    x_center = (x + width / 2) / img_w
    y_center = (y + height / 2) / img_h
    width /= img_w
    height /= img_h

    # 创建对应的YOLO标签文件
    label_filename = os.path.splitext(image_path)[0] + '.txt'
    label_file_path = os.path.join(output_path, label_filename)
    os.makedirs(os.path.dirname(label_file_path), exist_ok=True)

    with open(label_file_path, 'w') as f:
        f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

print("CUB200-2011数据集已成功转换为YOLO格式！")