在使用Ultralytics YOLO进行姿态估计模型训练时,需要遵循特定的数据集格式。每个图像文件对应一个文本文件,文件名与图像文件相同,扩展名为".txt"。文本文件中的每一行对应图像中的一个对象实例,包含对象的类别索引、中心坐标、宽度、高度以及关键点坐标等信息。
例如,对于二维姿态估计任务,标签格式如下:
<class-index> <x> <y> <width> <height> <px1> <py1> <px2> <py2> ... <pxn> <pyn>
对于三维姿态估计任务,标签格式稍有不同,每个关键点还包含可见性标志:
<class-index> <x> <y> <width> <height> <px1> <py1> <p1-visibility> <px2> <py2> <p2-visibility> ... <pxn> <pyn> <p2-visibility>
其中,<class-index>表示对象的类别索引,<x> <y> <width> <height>是边界框的坐标,<px1> <py1> <px2> <py2> ... <pxn> <pyn>是关键点的像素坐标。坐标之间以空格分隔。
Ultralytics框架使用YAML文件格式定义数据集和模型配置。以下是定义检测数据集的YAML格式示例:
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path:
../datasets/coco8-pose
train:
images/train
# train images (relative to 'path') 4 images
val:
images/val
# val images (relative to 'path') 4 images
test:
# test images (optional)
# Keypoints
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
]
# Classes dictionary
names:
0:
person
train和val字段指定包含训练和验证图像的目录路径,names是一个类别名称的字典。类别名称的顺序应与YOLO数据集文件中的类别索引顺序相匹配。如果关键点是对称的,则需要flip_idx,例如人体的左右侧或面部。
以下是使用Python CLI进行模型训练的示例:
from ultralytics import YOLO
# Load a model
model = YOLO("yolo11n-pose.pt")
# load a pretrained model (recommended for training)
# Train the model
results = model.train(data="coco8-pose.yaml", epochs=100, imgsz=640)
Ultralytics YOLO支持多种数据集,包括COCO-Pose、COCO8-Pose、Tiger-Pose和Hand Keypoints等,这些数据集都遵循Ultralytics YOLO格式,并可用于训练姿态估计模型。
如果有自己的数据集,并希望使用Ultralytics YOLO格式进行姿态估计模型训练,请确保数据集遵循上述"Ultralytics YOLO格式"。将注释转换为所需格式,并在YAML配置文件中指定路径、类别数量和类别名称。
Ultralytics提供了一个方便的转换工具,可以将COCO数据集格式的标签转换为YOLO格式:
from ultralytics.data.converter import convert_coco
convert_coco(labels_dir="path/to/coco/annotations/", use_keypoints=True)
这个转换工具可以用于将COCO数据集或任何COCO格式的数据集转换为Ultralytics YOLO格式。use_keypoints参数指定是否在转换的标签中包含关键点(用于姿态估计)。