[已解决]Python将COCO格式实例分割数据集转换为YOLO格式

直接根据代码更换自己成为自己的路径便可以转换

import json
import os
import numpy as np
from pycocotools import mask as coco_mask# 加载 COCO 数据集
coco_path = "/COCO2017/annotations/instances_train2017.json"
with open(coco_path, "r") as f:coco_data = json.load(f)# 创建保存 YOLO 标注的目录
output_dir = "/COCO2017/labels/train2017"
os.makedirs(output_dir, exist_ok=True)# 创建类别索引映射
category_id_to_index = {cat["id"]: idx for idx, cat in enumerate(coco_data["categories"])}# 创建图像ID到图像信息的映射
image_info_map = {img["id"]: img for img in coco_data["images"]}def rle_to_polygon(rle_segmentation, img_width, img_height):"""将RLE格式转换为多边形"""try:rle = coco_mask.frPyObjects(rle_segmentation, img_height, img_width)binary_mask = coco_mask.decode(rle)# 简单的轮廓提取（简化版）polygons = []# 这里需要更复杂的轮廓提取算法，暂时返回空列表return polygonsexcept:return []# 按图像ID组织标注
annotations_by_image = {}
for annotation in coco_data["annotations"]:image_id = annotation["image_id"]if image_id not in annotations_by_image:annotations_by_image[image_id] = []annotations_by_image[image_id].append(annotation)# 遍历每个图像的标注
for image_id, annotations in annotations_by_image.items():if image_id not in image_info_map:continueimage_info = image_info_map[image_id]image_width = image_info["width"]image_height = image_info["height"]image_filename = image_info["file_name"].split(".")[0]label_filename = os.path.join(output_dir, f"{image_filename}.txt")with open(label_filename, "w") as f:for annotation in annotations:category_id = annotation["category_id"]class_index = category_id_to_index[category_id]segmentation = annotation["segmentation"]polygons = []if isinstance(segmentation, dict):# RLE格式polygons = rle_to_polygon(segmentation, image_width, image_height)if not polygons:continueelif isinstance(segmentation, list) and segmentation:# 多边形格式polygons = segmentationelse:continuefor poly in polygons:if len(poly) < 6:continuenormalized_coords = []for i in range(0, len(poly), 2):if i + 1 < len(poly):x = max(0.0, min(1.0, poly[i] / image_width))y = max(0.0, min(1.0, poly[i + 1] / image_height))normalized_coords.extend([x, y])if len(normalized_coords) >= 6:coords_str = " ".join([f"{coord:.6f}" for coord in normalized_coords])f.write(f"{class_index} {coords_str}\n")print("转换完成！")

效果：