COCO格式数据集快转YOLO

import json
import os

def get_classes(coco_json_files):
    """提取所有分类"""
    classes = set()
    for coco_json in coco_json_files:
        with open(coco_json, 'r') as f:
            coco_data = json.load(f)
            for category in coco_data['categories']:
                classes.add(category['name'])
    return list(classes)

def convert_coco_to_yolo(coco_json, output_dir, category_mapping):
    """将 COCO 格式的标注转换为 YOLO 格式"""
    with open(coco_json, 'r') as f:
        coco_data = json.load(f)

    image_info = {img['id']: img for img in coco_data['images']}
    annotations_count = 0

    for ann in coco_data['annotations']:
        img_id = ann['image_id']
        img_info = image_info[img_id]
        img_width = img_info['width']
        img_height = img_info['height']

        category_id = category_mapping[coco_data['categories'][ann['category_id'] - 1]['name']]
        bbox = ann['bbox']
        x_min, y_min, width, height = bbox

        x_center = (x_min + width / 2) / img_width
        y_center = (y_min + height / 2) / img_height
        width /= img_width
        height /= img_height

        yolo_label = f"{category_id} {x_center} {y_center} {width} {height}"
        yolo_file = os.path.join(output_dir, f"{img_info['file_name'].replace('.jpg', '.txt')}")

        os.makedirs(os.path.dirname(yolo_file), exist_ok=True)
        with open(yolo_file, 'a') as f:
            f.write(yolo_label + '\n')

        annotations_count += 1

    return annotations_count

def process_directory(input_dir, output_dir):
    """递归处理目录"""
    coco_json_files = []
    for root, _, files in os.walk(input_dir):
        for file in files:
            if file.endswith('.json'):
                coco_json_files.append(os.path.join(root, file))

    # 提取所有分类
    all_classes = get_classes(coco_json_files)
    category_mapping = {name: idx for idx, name in enumerate(sorted(all_classes))}

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 将分类写入 classes.txt
    classes_txt_path = os.path.join(output_dir, 'classes.txt')
    with open(classes_txt_path, 'w') as f:
        for category_name in sorted(all_classes):
            f.write(category_name + '\n')

    # 处理每个 COCO JSON 文件，并打印标注数
    total_annotations = 0
    for coco_json in coco_json_files:
        annotations_count = convert_coco_to_yolo(coco_json, output_dir, category_mapping)
        total_annotations += annotations_count
        print(f"Processed {coco_json}: {annotations_count} annotations")

    print(f"Total annotations processed: {total_annotations}")
    print(f"Total classes: {len(all_classes)}")

# 使用示例
process_directory('C:\\Users\\合格的牛马\\Desktop\\test', 'C:\\Users\\合格的牛马\\Desktop\\test\\demo')