COCO格式的数据集标注文件(JSON格式)

转换为YOLO格式的数据集标注文件(TXT格式)

在这里插入图片描述

使用方法:修改coco的位置,以及保存的文件夹位置,然后直接运行下方命令
有问题问AI
python cocozhuan3yolo.py --coco "C:\Users\AnK\Desktop\Seaclear Marine Debris Dataset\dataset.json" --output "C:\Users\AnK\Desktop\output"
import json
import os
from collections import defaultdict
import argparse
#使用方法:修改coco的位置,以及保存的文件夹位置,然后直接运行下方命令
#有问题问AI
#python cocozhuan3yolo.py --coco "C:\Users\AnK\Desktop\Seaclear Marine Debris Dataset\dataset.json" --output "C:\Users\AnK\Desktop\output"

def convert_coco_to_yolo(coco_json_path, output_dir):
    # 读取COCO格式的JSON文件
    with open(coco_json_path, 'r') as f:
        data = json.load(f)

    # 解析类别信息,并建立 COCO 类别 id 到 YOLO 类别索引的映射(从 0 开始)
    categories = data["categories"]
    # 按照原有的 id 顺序排序(也可以根据 name 字母排序,前提是一致)
    categories = sorted(categories, key=lambda x: x["id"])
    # 构造映射字典:COCO的类别ID -> YOLO中的类别索引
    cat_id_to_yolo = {}
    for idx, cat in enumerate(categories):
        cat_id_to_yolo[cat["id"]] = idx

    # 建立图像id到图像信息(如宽高、文件名)的字典
    images_info = {}
    for img in data["images"]:
        images_info[img["id"]] = img

    # 分组收集每张图片的标注信息
    ann_by_img = defaultdict(list)
    if "annotations" in data:
        annotations = data["annotations"]
        for ann in annotations:
            image_id = ann["image_id"]
            ann_by_img[image_id].append(ann)
    else:
        print("未在JSON文件中找到annotations字段!")
        return

    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 遍历每个图像的标注,生成YOLO格式的txt文件
    for image_id, annotations in ann_by_img.items():
        # 获取当前图像信息及其尺寸
        img_info = images_info[image_id]
        img_width = img_info["width"]
        img_height = img_info["height"]
        # 保留文件名(不含扩展名)作为标注txt文件名字
        base_name = os.path.splitext(img_info["file_name"])[0]
        txt_path = os.path.join(output_dir, base_name + ".txt")

        yolo_lines = []
        for ann in annotations:
            # 获取COCO类别和转换为YOLO的类别索引
            coco_cat_id = ann["category_id"]
            yolo_cat_id = cat_id_to_yolo[coco_cat_id]
            # COCO bbox 格式:[x_min, y_min, width, height]
            bbox = ann["bbox"]
            x_min, y_min, bbox_width, bbox_height = bbox
            # 计算中心点坐标
            x_center = x_min + bbox_width / 2.0
            y_center = y_min + bbox_height / 2.0
            # 归一化坐标和尺寸
            x_center_norm = x_center / img_width
            y_center_norm = y_center / img_height
            width_norm = bbox_width / img_width
            height_norm = bbox_height / img_height

            # 生成一行YOLO格式标注,保留6位小数
            line = f"{yolo_cat_id} {x_center_norm:.6f} {y_center_norm:.6f} {width_norm:.6f} {height_norm:.6f}"
            yolo_lines.append(line)

        # 将标注写入对应的txt文件中
        with open(txt_path, "w") as f:
            f.write("\n".join(yolo_lines))

    # 生成类别名称文件(例如:classes.txt),每行一个类别,顺序与 YOLO 的索引一致
    classes_path = os.path.join(output_dir, "classes.txt")
    with open(classes_path, "w") as f:
        for cat in categories:
            f.write(cat["name"] + "\n")

    print(f"转换完成!YOLO格式的标注文件和类别文件已经保存在:{output_dir}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="将COCO格式数据集转换为YOLO格式数据集")
    parser.add_argument("--coco", type=str, required=True, help="COCO JSON文件路径,例如:dataset.json")
    parser.add_argument("--output", type=str, required=True, help="输出文件夹路径,用于保存YOLO格式的标注文件和类别文件")
    args = parser.parse_args()

    convert_coco_to_yolo(args.coco, args.output)

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐