yolov8训练pose（人体姿态）数据集

目录

一、源码

二、数据集制作和格式转换

2.1 数据集制作

2.2 数据集格式转换

2.2.1 需要修改的地方

三、训练

---

一、源码

官方源码

二、数据集制作和格式转换

2.1 数据集制作

打开labelme，加载图片

先用矩形框标注人

再用point点标注人体关节点

***  可以只标注看得见的关节点***看不见的可以不用标注

这是关节点图

标注完毕后的数据集格式，图片和json放一起。数据集的存放位置如图

2.2 数据集格式转换

在项目跟目录下新建一个 label2yolo.py 文件，复制以下代码

import os
import json
import shutil
import random
import argparse
from tqdm import tqdm
from collections import Counter
import yaml

# you need to modify keypoint_class and bbox_class to match the yaml file   关键点名字、数量
keypoint_class = ['0','1','2','3','4','5','6','7','8','9','10','11','12','13','14','15','16']
# 物体的类名
bbox_class = {
    'person': 0
}


def get_classes(json_dir):
    '''
    统计路径下 JSON 文件里的各类别标签数量
    '''
    names = []
    json_files = [os.path.join(json_dir, f) for f in os.listdir(json_dir) if f.endswith('.json')]

    for json_path in json_files:
        with open(json_path, 'r') as f:
            data = json.load(f)
            for shape in data['shapes']:
                name = shape['label']
                if name in bbox_class:
                    names.append(name)
                else:
                    continue

    result = Counter(names)
    return result


def process_single_json(labelme_path, save_folder):
    with open(labelme_path, 'r', encoding='utf-8') as f:
        labelme = json.load(f)

    img_width = labelme['imageWidth']
    img_height = labelme['imageHeight']

    suffix = labelme_path.split('.')[-2]
    yolo_txt_path = suffix + '.txt'

    with open(yolo_txt_path, 'w', encoding='utf-8') as f:

        for each_ann in labelme['shapes']:

            if each_ann['shape_type'] == 'rectangle':

                yolo_str = ''

                bbox_class_id = bbox_class[each_ann['label']]
                yolo_str += '{} '.format(bbox_class_id)

                bbox_top_left_x = int(min(each_ann['points'][0][0], each_ann['points'][1][0]))
                bbox_bottom_right_x = int(max(each_ann['points'][0][0], each_ann['points'][1][0]))
                bbox_top_left_y = int(min(each_ann['points'][0][1], each_ann['points'][1][1]))
                bbox_bottom_right_y = int(max(each_ann['points'][0][1], each_ann['points'][1][1]))

                bbox_center_x = int((bbox_top_left_x + bbox_bottom_right_x) / 2)
                bbox_center_y = int((bbox_top_left_y + bbox_bottom_right_y) / 2)

                bbox_width = bbox_bottom_right_x - bbox_top_left_x

                bbox_height = bbox_bottom_right_y - bbox_top_left_y

                bbox_center_x_norm = bbox_center_x / img_width
                bbox_center_y_norm = bbox_center_y / img_height

                bbox_width_norm = bbox_width / img_width

                bbox_height_norm = bbox_height / img_height

                yolo_str += '{:.5f} {:.5f} {:.5f} {:.5f} '.format(bbox_center_x_norm, bbox_center_y_norm,
                                                                  bbox_width_norm, bbox_height_norm)

                bbox_keypoints_dict = {}
                for each_ann in labelme['shapes']:
                    if each_ann['shape_type'] == 'point':
                        x = int(each_ann['points'][0][0])
                        y = int(each_ann['points'][0][1])
                        label = each_ann['label']
                        if (x > bbox_top_left_x) & (x < bbox_bottom_right_x) & (y < bbox_bottom_right_y) & (
                                y > bbox_top_left_y):
                            bbox_keypoints_dict[label] = [x, y]

                for each_class in keypoint_class:
                    if each_class in bbox_keypoints_dict:
                        keypoint_x_norm = bbox_keypoints_dict[each_class][0] / img_width
                        keypoint_y_norm = bbox_keypoints_dict[each_class][1] / img_height
                        yolo_str += '{:.5f} {:.5f} {} '.format(keypoint_x_norm, keypoint_y_norm, 2)
                    else:
                        yolo_str += '0 0 0 '
                f.write(yolo_str + '\n')

    shutil.move(yolo_txt_path, save_folder)


def json_get_class(dataset_path, output_path):
    # 统计路径下 JSON 文件里的各类别标签数量
    obj_classes = get_classes(dataset_path)
    classes = list(obj_classes.keys())

    # 编写yaml文件
    classes_txt = {i: classes[i] for i in range(len(classes))}  # 标签类别
    data = {
        'path': os.path.join(os.getcwd(), output_path),
        'train': "images/train",
        'val': "images/val",
        'names': classes_txt,
        'nc': len(classes)
    }

    with open(output_path + '/pose.yaml', 'w', encoding="utf-8") as file:
        yaml.dump(data, file)
    print("标签：", dict(obj_classes))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="labelme annotation to yolo dataset.")

    # input params
    parser.add_argument('--data_path', type=str, default=r"E:\python\yolov8\ultralytics_train_7_30\datasets\pose\image", help='dataset folder path to convert: ')
    parser.add_argument('--output_path', type=str, default=r"E:\python\yolov8\ultralytics_train_7_30\datasets\pose\myyolo", help='output folder path to save: ')
    args = parser.parse_args()
    dataset_path = args.data_path
    output_path = args.output_path

    print("dataset path: ", dataset_path)
    print("output_path path: ", output_path)

    if not os.path.isdir(output_path):
        os.makedirs(output_path)
    if not os.path.isdir(os.path.join(output_path, 'images', 'train')):
        os.makedirs(os.path.join(output_path, 'images', 'train'))
    if not os.path.isdir(os.path.join(output_path, 'images', 'val')):
        os.makedirs(os.path.join(output_path, 'images', 'val'))
    if not os.path.isdir(os.path.join(output_path, 'labels', 'train')):
        os.makedirs(os.path.join(output_path, 'labels', 'train'))
    if not os.path.isdir(os.path.join(output_path, 'labels', 'val')):
        os.makedirs(os.path.join(output_path, 'labels', 'val'))

    # select train and val dataset
    os.chdir(dataset_path)

    json_files = list(filter(lambda x: '.json' in x, os.listdir()))
    json_ext = json_files[0].split(".")[-1]

    img_files = list(filter(lambda x: '.json' not in x, os.listdir()))
    img_ext = img_files[0].split(".")[-1]

    files_without_ext = [file.split(".")[0] for file in json_files]

    frac = 0.2
    random.seed(123)
    val_num = int(len(files_without_ext) * frac)
    train_files = files_without_ext[val_num:]
    val_files = files_without_ext[:val_num]

    print("Total: ", len(files_without_ext))
    print("Train: ", len(train_files))
    print("Value: ", len(val_files))

    # copy to output file
    for train_file in tqdm(train_files):
        img_train_file = train_file + '.' + img_ext
        json_train_file = train_file + '.' + json_ext
        shutil.copy(img_train_file, os.path.join(output_path, 'images', 'train'))
        process_single_json(json_train_file, os.path.join(output_path, 'labels', 'train'))

    for val_file in tqdm(val_files):
        img_val_file = val_file + '.' + img_ext
        json_val_file = val_file + '.' + json_ext
        shutil.copy(img_val_file, os.path.join(output_path, 'images', 'val'))
        process_single_json(json_val_file, os.path.join(output_path, 'labels', 'val'))

    json_get_class(dataset_path, output_path)

    print("Successful!")

2.2.1 需要修改的地方

如果你是训练非人体的目标，需要修改关键点的名字，数量，以及目标的类名；如果你训练的是人体姿态就不需要修改

修改数据集路径（*** 需要用绝对路径 ***）

运行后得到分割好的数据集，以及用于训练的yaml文件。

打开yaml文件，将以下内容复制到进去

kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
flip_idx: [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]

最终长这样

至此，数据集、yaml配置文件处理完毕

三、训练

在ultralytics/models/yolo/pose路径下新建一个pose_train_my.py，将以下内容复制进去。需提前在官网下载模型

from ultralytics import YOLO

# Load a model
model = YOLO("yolov8n-pose.pt")  # load a pretrained model (recommended for training)


# Train the model
results = model.train(task='pose',
                      data="datasets/pose/myyolo/pose.yaml",
                      model='yolov8n-pose.pt',
                      epochs=300,
                      imgsz=640,
                      batch=16,
                      workers=0,
                      save_period=10, # 多少轮保存一个模型（-1 不保存）
                      resume=False)

运行pose_train_my.py开始训练，参数可以用我设置的，也可以自己修改

训练完后模型保存在同目录下的runs文件夹下