python生成 coco 关键点标注数据

import json
import os
import cv2  # 用于读取图像尺寸

# 假设你有以下结构
# your_data = [
#   {"file_name": "000000000001.jpg", "keypoints": [...], "bbox": [...]},
#   ...
# ]

your_data = [...]  # 你的数据列表，每项包含 file_name、keypoints 和 bbox

output = {
    "images": [],
    "annotations": [],
    "categories": [
        {
            "id": 1,
            "name": "person",
            "supercategory": "person",
            "keypoints": [
                "nose", "left_eye", "right_eye", "left_ear", "right_ear",
                "left_shoulder", "right_shoulder", "left_elbow", "right_elbow",
                "left_wrist", "right_wrist", "left_hip", "right_hip",
                "left_knee", "right_knee", "left_ankle", "right_ankle"
            ],
            "skeleton": [
                [16,14],[14,12],[17,15],[15,13],
                [12,13],[6,12],[7,13],[6,7],[6,8],
                [7,9],[8,10],[9,11],[2,3],[1,2],[1,3],
                [2,4],[3,5],[4,6],[5,7]
            ]
        }
    ]
}

for idx, item in enumerate(your_data):
    file_path = os.path.join("your/images/dir", item["file_name"])
    
    # 读取图片尺寸
    img = cv2.imread(file_path)
    height, width = img.shape[:2]

    # 图像元数据
    image_id = idx + 1
    output["images"].append({
        "id": image_id,
        "file_name": item["file_name"],
        "width": width,
        "height": height
    })

    # 注释数据
    keypoints = item["keypoints"]  # 长度应为 51 (17 x 3)
    num_keypoints = sum(1 for i in range(2, len(keypoints), 3) if keypoints[i] > 0)

    output["annotations"].append({
        "id": image_id,
        "image_id": image_id,
        "category_id": 1,
        "keypoints": keypoints,
        "num_keypoints": num_keypoints,
        "bbox": item["bbox"],  # COCO格式是[x, y, width, height]
        "area": item["bbox"][2] * item["bbox"][3],
        "iscrowd": 0
    })

# 保存为 JSON 文件
with open("person_keypoints_val2014.json", "w") as f:
    json.dump(output, f, indent=2)