已有YOLO格式数据集一键转换为COCO格式

当我们已经有了YOLO格式的训练集和验证集的txt文件时，这两个里面均存储了图片地址，图片地址也对应于标签地址：

图片地址：

bash 复制代码

/data/YOLO/images/1.jpg
/data/YOLO/images/2.jpg

标签地址：

bash 复制代码

/data/YOLO/labels/1.txt
/data/YOLO/labels/2.txt

path/to/coco/

annotations/ # annotation json files

train/ # train images

val/ # val images

python 复制代码

import os
import json
import shutil
from PIL import Image
from tqdm import tqdm

def create_coco_structure(base_dir):
    """
    创建 COCO 数据集所需的基本目录结构。
    """
    train_dir = os.path.join(base_dir, 'train')
    val_dir = os.path.join(base_dir, 'val')
    annotations_dir = os.path.join(base_dir, 'annotations')

    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(annotations_dir, exist_ok=True)

    print(f"成功创建目录结构 at '{base_dir}'")
    return train_dir, val_dir, annotations_dir

def yolo_to_coco_converter(file_list_path, output_image_dir):
    """
    将 YOLO 格式的数据集转换为 COCO 格式。

    Args:
        file_list_path (str): 包含图片绝对路径的 .txt 文件 (e.g., train.txt)。
        output_image_dir (str): 转换后图片存放的目录。

    Returns:
        tuple: 包含 COCO 格式的 images 和 annotations 列表。
    """
    images_info = []
    annotations_info = []
    image_id_counter = 0
    annotation_id_counter = 0

    with open(file_list_path, 'r') as f:
        image_paths = [line.strip() for line in f.readlines() if line.strip()]

    print(f"开始处理 {os.path.basename(file_list_path)} 中的 {len(image_paths)} 张图片...")

    for img_path in tqdm(image_paths):
        if not os.path.exists(img_path):
            print(f"警告: 图片路径不存在，已跳过: {img_path}")
            continue
        
        # 复制图片到目标文件夹
        img_filename = os.path.basename(img_path)
        dest_img_path = os.path.join(output_image_dir, img_filename)
        shutil.copy(img_path, dest_img_path)

        # 获取图片尺寸
        try:
            with Image.open(img_path) as img:
                width, height = img.size
        except IOError:
            print(f"警告: 无法打开或读取图片，已跳过: {img_path}")
            continue

        # 添加图片信息到 COCO images 列表
        image_info = {
            "id": image_id_counter,
            "file_name": img_filename,
            "width": width,
            "height": height
        }
        images_info.append(image_info)

        # 4. 处理对应的标签文件
        label_path = img_path.replace('/images/', '/labels/').replace('.jpg', '.txt').replace('.png', '.txt')
        
        if os.path.exists(label_path):
            with open(label_path, 'r') as lf:
                for line in lf.readlines():
                    parts = line.strip().split()
                    if len(parts) != 5:
                        continue

                    class_id, x_center, y_center, bbox_width, bbox_height = map(float, parts)
                    
                    # YOLO 归一化坐标 -> COCO 绝对坐标 [x_min, y_min, width, height]
                    abs_bbox_width = bbox_width * width
                    abs_bbox_height = bbox_height * height
                    x_min = (x_center * width) - (abs_bbox_width / 2)
                    y_min = (y_center * height) - (abs_bbox_height / 2)

                    annotation = {
                        "id": annotation_id_counter,
                        "image_id": image_id_counter,
                        "category_id": int(class_id),
                        "bbox": [x_min, y_min, abs_bbox_width, abs_bbox_height],
                        "area": abs_bbox_width * abs_bbox_height,
                        "iscrowd": 0
                    }
                    annotations_info.append(annotation)
                    annotation_id_counter += 1
        
        image_id_counter += 1
    
    return images_info, annotations_info

def main():
    # --------------------------------------------------------------------------------------------------------------- #
    # 重要: 请修改 CATEGORIES 列表以匹配数据集类别。
    # 顺序必须与 YOLO .txt 文件中的类别 ID 相同。
    # 例如, 如果 'person' 是类别 0, 'car' 是类别 1, 那么:
    # CATEGORIES = ["person", "car"]
    # --------------------------------------------------------------------------------------------------------------- #
    CATEGORIES = ["car", "bus", "person"] #根据自己的类别添加

    TRAIN_TXT_PATH = '/data/YOLO/train.txt'
    VAL_TXT_PATH = '/data/YOLO/val.txt'
    OUTPUT_BASE_DIR = '/data/coco'

    train_image_dir, val_image_dir, annotations_dir = create_coco_structure(OUTPUT_BASE_DIR)
    coco_categories = [{"id": i, "name": name, "supercategory": "object"} for i, name in enumerate(CATEGORIES)]
    train_images, train_annotations = yolo_to_coco_converter(TRAIN_TXT_PATH, train_image_dir)
    train_coco_format = {
        "info": {"description": "Train dataset in COCO format"},
        "licenses": [],
        "images": train_images,
        "annotations": train_annotations,
        "categories": coco_categories
    }
    train_json_path = os.path.join(annotations_dir, 'train.json')
    with open(train_json_path, 'w') as f:
        json.dump(train_coco_format, f, indent=4)
    print(f"成功生成训练集标注文件: {train_json_path}")

    # --- 处理验证集 ---
    val_images, val_annotations = yolo_to_coco_converter(VAL_TXT_PATH, val_image_dir)
    val_coco_format = {
        "info": {"description": "Validation dataset in COCO format"},
        "licenses": [],
        "images": val_images,
        "annotations": val_annotations,
        "categories": coco_categories
    }
    val_json_path = os.path.join(annotations_dir, 'val.json')
    with open(val_json_path, 'w') as f:
        json.dump(val_coco_format, f, indent=4)
    print(f"成功生成验证集标注文件: {val_json_path}")

    print("\n转换完成！")

if __name__ == '__main__':
    main()