【目标检测】COCO格式数据集转换为YOLO格式

COCO2YOLO， 修改category_mapping:

import json
import os

def coco_to_yolo(coco_annotation_file, output_dir):
    # Load COCO annotations
    with open(coco_annotation_file, 'r') as f:
        coco_data = json.load(f)

    # Create a dictionary to map category_id to class_id
    # Reorder categories to match the desired YOLO indices
    category_mapping = {
        1: 0,  # class 1
        2: 1,  # class 2
        3: 2   # class 3
    }
    
    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)
    
    # Iterate through images
    for img in coco_data['images']:
        img_id = img['id']
        img_width = img['width']
        img_height = img['height']
        annotations = [ann for ann in coco_data['annotations'] if ann['image_id'] == img_id]
        
        # Prepare the YOLO annotation content
        yolo_annotations = []
        for ann in annotations:
            x, y, w, h = ann['bbox']
            x_center = (x + w / 2) / img_width
            y_center = (y + h / 2) / img_height
            w = w / img_width
            h = h / img_height
            category_id = ann['category_id']
            class_id = category_mapping[category_id]
            yolo_annotations.append(f"{class_id} {x_center} {y_center} {w} {h}")
        
        # Save YOLO annotations to a file
        output_file = os.path.join(output_dir, f"{img['file_name'].split('.')[0]}.txt")
        with open(output_file, 'w') as f:
            f.write('\n'.join(yolo_annotations))

# Usage example
coco_annotation_file = '/data/instances_val2017.json'
output_dir = '/data/yolo_annotations'
coco_to_yolo(coco_annotation_file, output_dir)