coco数据集格式转yolo数据格式

一、coco数据集是什么？

COCO（Common Objects in Context）是一个广泛使用的目标检测和分割数据集，而YOLO（You Only Look Once）是一种流行的实时目标检测算法。

首先，导入了必要的库，包括json和os。然后，定义了一个名为convert_coco_to_yolo的函数，用于将COCO数据转换为YOLO格式。该函数接受一个COCO数据的字典作为输入，并返回转换后的YOLO数据和类别字典。

在函数内部，提取了COCO数据中的图像列表、注释和类别信息。然后，创建了一个空的YOLO数据列表和一个用于存储类别名和序号映射关系的字典。

接下来，遍历图像列表，并获取图像的ID、文件名、宽度和高度。然后，根据图像ID筛选出与该图像相关的注释，并将它们存储在image_annotations列表中。

然后，遍历image_annotations列表，并提取注释的类别ID。通过在类别列表中查找与类别ID匹配的类别，可以获取类别的名称。如果找不到匹配的类别，则跳过该注释。

接下来，提取注释的边界框信息，并计算出边界框的中心坐标、归一化后的坐标和大小。将这些信息存储在yolo_annotations列表中。

如果存在YOLO注释，则按照类别序号对注释进行排序，并将图像的相关信息和注释存储在yolo_data列表中。

最后，返回转换后的YOLO数据和类别字典。

接下来，定义了几个变量，包括COCO数据文件的路径、文件名和保存目录的路径。然后，使用os.path.join函数构建了完整的文件路径。

通过检查文件是否存在，可以确保文件路径是有效的。如果文件存在，打开文件并加载其中的JSON数据。然后，调用convert_coco_to_yolo函数将COCO数据转换为YOLO数据。

使用os.makedirs函数创建保存目录，如果目录已经存在，则不进行任何操作。

接下来，生成了一个class.txt文件，用于存储类别的名称。遍历类别字典，并按照类别序号对字典进行排序。然后，将类别的名称写入文件中。

最后，遍历YOLO数据列表，并根据图像的文件名创建一个对应的.txt文件。然后，遍历注释列表，并将注释的类别和其他信息写入文件中。

最后，打印转换完成的消息，并显示保存目录的路径。如果文件不存在，则打印文件不存在的消息。

通过运行这个程序，可以将COCO数据集格式转换为YOLO格式，并将转换后的数据保存到指定的目录中。

这篇博客介绍了一个将COCO数据集格式转换为YOLO格式的Python程序。通过这个程序，可以轻松地处理COCO数据，并将其转换为适用于YOLO算法的格式。希望这个程序对你有所帮助！

二、完整代码

import json
import os

def convert_coco_to_yolo(coco_data):
    image_list = coco_data['images']
    annotations = coco_data['annotations']
    categories = coco_data['categories']

    yolo_data = []
    category_dict = {}  # 用于存储类别名和对应的序号
    for i, category in enumerate(categories):  # 构建类别名和序号的映射关系
        category_dict[category['name']] = i

    for image in image_list:
        image_id = image['id']
        file_name = image['file_name']
        width = image['width']
        height = image['height']

        image_annotations = [ann for ann in annotations if ann['image_id'] == image_id]
        yolo_annotations = []
        for ann in image_annotations:
            category_id = ann['category_id']
            category = next((cat for cat in categories if cat['id'] == category_id), None)
            if category is None:
                continue

            bbox = ann['bbox']
            x, y, w, h = bbox
            x_center = x + w / 2
            y_center = y + h / 2
            normalized_x_center = x_center / width
            normalized_y_center = y_center / height
            normalized_width = w / width
            normalized_height = h / height

            yolo_annotations.append({
                'category': category_dict[category['name']],  # 使用类别序号
                'x_center': normalized_x_center,
                'y_center': normalized_y_center,
                'width': normalized_width,
                'height': normalized_height
            })

        if yolo_annotations:
            yolo_annotations.sort(key=lambda x: x['category'])  # 按类别序号排序
            yolo_data.append({
                'file_name': file_name,
                'width': width,
                'height': height,
                'annotations': yolo_annotations
            })

    return yolo_data, category_dict

path = '/codeyard/yolov5_6.0/data_MS_1001labels'  # 修改为包含 via_export_coco.json 文件的目录路径
file_name = 'coco_filtered.json'  # 文件名
save_dir = '/codeyard/yolov5_6.0/data_MS_1001labels/labels'  # 保存目录
file_path = os.path.join(path, file_name)  # 完整文件路径

if os.path.isfile(file_path):  # 检查文件是否存在
    with open(file_path, 'r', encoding='utf-8') as load_f:
        load_dict = json.load(load_f)
        yolo_data, category_dict = convert_coco_to_yolo(load_dict)

    os.makedirs(save_dir, exist_ok=True)  # 创建保存目录

    # 生成 class.txt 文件
    class_file_path = os.path.join(save_dir, 'classes.txt')
    with open(class_file_path, 'w', encoding='utf-8') as class_f:
        for category_name, category_index in sorted(category_dict.items(), key=lambda x: x[1]):
            class_f.write(f"{category_name}\n")

    for data in yolo_data:
        file_name = os.path.basename(data['file_name'])  # 提取文件名部分
        width = data['width']
        height = data['height']
        annotations = data['annotations']

        txt_file_path = os.path.join(save_dir, os.path.splitext(file_name)[0] + '.txt')
        with open(txt_file_path, 'w', encoding='utf-8') as save_f:
            for annotation in annotations:
                category = annotation['category']
                x_center = annotation['x_center']
                y_center = annotation['y_center']
                box_width = annotation['width']
                box_height = annotation['height']

                line = f"{category} {x_center:.6f} {y_center:.6f} {box_width:.6f} {box_height:.6f}\n"
                save_f.write(line)

    print("转换完成，保存到:", save_dir)
else:
    print("文件不存在:", file_path)