coco数据集转换SAM2格式

coco是一个大json汇总了所有train的标签
SAM2训练一张图对应一个json标签

import json
import os
from pycocotools import mask as mask_utils
import numpy as np
import cv2def poly2mask(points, width, height):points_array = np.array(points, dtype=np.int32).reshape(-1, 2)mask = np.zeros((height, width), dtype=np.uint8)  # 注意顺序是(height, width)cv2.fillPoly(mask, [points_array], 255)  # 填充多边形区域为255return mask2rle(mask)def mask2rle(mask):"""将二值化掩码转换为RLE编码"""rle = mask_utils.encode(np.asfortranarray(mask))  # 使用pycocotools进行RLE编码rle['counts'] = rle['counts'].decode('utf-8')  # 将bytes转换为字符串return rle# 读取COCO格式JSON文件
with open('/home//Datasets/coco12cup/train.json', 'r') as f:coco_data = json.load(f)# 创建目标文件夹（如果不存在）
output_dir = '/home//Datasets/coco12cup/train'
os.makedirs(output_dir, exist_ok=True)# 遍历图像信息
for image in coco_data['images']:image_id = image['id']height = image['height']width = image['width']file_name = image['file_name']# 筛选对应图像的标注信息annotations = [anno for anno in coco_data['annotations'] if anno['image_id'] == image_id]# 转换标注格式为SAM2格式sam2_annotations = []for anno in annotations:# 检查segmentation格式segmentation = anno['segmentation']if isinstance(segmentation, list):  # 多边形格式segmentation_rle = poly2mask(segmentation[0], width, height)  # 多边形可能有多个，需要选第一个或合并else:  # 如果是RLE格式segmentation_rle = segmentationsam2_anno = {'area': anno['area'],'bbox': anno['bbox'],'id': anno['id'],'segmentation': segmentation_rle,}sam2_annotations.append(sam2_anno)# 创建SAM2格式的JSON数据sam2_data = {'annotations': sam2_annotations,'image': {'date_captured': '20241210',  # 根据实际情况修改'file_name': file_name,'height': height,'image_id': image_id,'license': 1,  # 根据实际情况修改'width': width}}# 写入单个JSON文件（使用 file_name 替换原本的 id 作为文件名）output_name = os.path.splitext(file_name)[0] + '.json'  # 替换扩展名为 .jsonoutput_path = os.path.join(output_dir, output_name)with open(output_path, 'w') as f:json.dump(sam2_data, f)