YOLOv8关键点检测数据转换及划分脚本

原创

已于 2025-03-14 09:29:44 修改 · 312 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #人工智能 #python

于 2024-04-02 14:39:37 首次发布

本文介绍了如何使用Python脚本对图像和标签文件进行组织，创建训练集和验证集，以支持YOLO（YouOnlyLookOnce）等深度学习对象检测模型的训练。函数`create_folders`用于创建所需的文件夹结构，`split_dataset`则根据给定的训练比例划分数据集。

在计算机视觉任务中，关键点检测（如人体姿态估计、面部关键点检测等）需要将标注数据转换为模型训练所需的格式。本文将详细介绍如何将JSON格式的关键点数据集转换为TXT格式，并合理划分训练集、验证集和测试集。

一、JSON转TXT格式详解

1.JSON格式解析

以常见的COCO关键点标注格式为例，JSON文件通常包含以下字段：

images: 图像信息（文件名、尺寸等）
annotations: 标注信息（边界框、关键点坐标及可见性）
categories: 类别定义（如关键点名称、骨架连接关系）

import os
import cv2
import numpy as np
import glob
import json
import tqdm

# 物体类别
class_list = ['person']
# 关键点的顺序
keypoint_list = ['1', '2']

# 新的txt文件保存路径（假设为当前目录下的"output_txts/"）
output_folder = "./Yolo_Data/"


def create_output_folder_if_not_exists(folder_path):
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)


def json_to_yolo(img_data, json_data):
    h, w = img_data.shape[:2]

    rectangles = {
   
   }
    # 遍历初始化
    for shape in json_data["shapes"]:
        label = shape["label"]  # pen, head, tail
        points = shape["points"]  # x,y coordinates
        shape_type = shape["shape_type"]

        # 只处理矩形
        if shape_type == "rectangle":
            if label in class_list:
                rect_id = len(rectangles)
                rectangles[rect_id] = {
   
   
                    "label": label,
                    "rect": points[0] + points[1],  # Rectangle [x1, y1, x2, y2]
                    "keypoints_list": []
                }

    # 遍历更新，将点加入对应的目标框中
    for shape in json_data["shapes"]:
        label = shape["label"]
        points = shape["points"]
        # 如果匹配到了对应的keypoint