基于Yolov8-pose自制手部姿态数据集
下载链接在我的资源里，支持开源。
目前网上有许多手部检测或手部关键点检测的数据集，但是同时具有bounding box和 keypoints 的数据集很少。
通过网上翻找，发现貌似只有22年新出的数据集Hagrid v1符合这个条件。
注意在官网github上Hagrid具有v1和v2两个版本，其中只有v1版本的标注中同时具有关键点和手的bbox，我们将Hagrid v1版本的项目git到本地。
由于这个数据集中图片的分辨率很高，因此非常庞大。我们选择下载他的mini版本，也就是这里的链接，下载Subsample中的图像数据集images和标注annotations。
By the way，如果用全部数据的话用本文的方法也是可以走的通的，因为数据格式一样，无非就是改改路径。实测只用这些子类的100张数据表现的效果不是很好，建议用全部数据好一些。全部数据总共走下来有414819个label，数据十分庞大。
注意，yolov8-pose的标注格式要求为
#以6个关键点为例
类别0 + bbox [cx cy w h] + 关键点 [k1, k2 ,k3 ,k4 ,k5 ,k6 ]
拿到标注和图片以后，我们使用Hagrid中自带的脚本converters.hagrid_to_yolo.py来将我们的数据集转化成yolov8-pose的格式，代码如下：
#hagrid_to_yolo.py
import argparse
import json
import logging
import os
from typing import Tuple, Union
import numpy as np
import pandas as pd
from omegaconf import DictConfig, ListConfig, OmegaConf
from tqdm import tqdm
from constants import IMAGES
import shutil
logging.getLogger().setLevel(logging.INFO)
tqdm.pandas()
def xywh_to_cxcywh(boxes):
"""
Convert xywh to cxcywh
Parameters
----------
boxes: list
list of bboxes
Returns
-------
list
"""
boxes = np.array(boxes)
return np.concatenate([boxes[..., :2] + boxes[..., 2:] / 2, boxes[..., 2:]], len(boxes.shape) - 1)
def xywh_to_xyxy(boxes):
"""
Convert xywh to xyxy
Parameters
----------
boxes: list
list of bboxes
Returns
-------
list
"""
boxes = np.array(boxes)
return np.concatenate([boxes[..., :2], boxes[..., :2] + boxes[..., 2:]], len(boxes.shape) - 1)
def get_files_from_dir(pth: str, extns: Tuple) -> list:
"""
Get files from directory
Parameters
----------
pth: str
path to directory
extns: Tuple
extensions of files
Returns
-------
list
"""
if not os.path.exists(pth):
logging.error(f"Dataset directory doesn't exist {
pth}")
return []
files = [f for f in os.listdir(pth) if f.endswith(extns)]
return files
def get_dataframe(conf: Union[DictConfig, ListConfig], phase: str) -> pd.DataFrame:
"""
Get dataframe with annotations
Parameters
----------
conf: Union[DictConfig, ListConfig]
config
phase: str
phase of dataset
Returns
-------
pd.DataFrame
"""
dataset_annotations = conf.dataset.dataset_annotations
dataset_folder = conf.dataset.dataset_folder
targets = conf.dataset.targets
annotations_all = None
exists_images = []
for target in tqdm(targets):
target_json = os.path.join(dataset_annotations, f"ann_{
phase}", f"{
target}.json")
if os.path.exists(target_json):
json_annotation = json.load(open(os.path.join(target_json)))
json_annotation = [
dict(annotation, **{
"name": f"{
name}.jpg"})
for name, annotation in zip(json_annotation, json_annotation.values())
]
annotation = pd.DataFrame(json_annotation)
annotation["target"] = target
annotations_all = pd.concat([annotations_all, annotation], ignore_index=True)
exists_images.extend(get_files_from_dir(os.path.join(dataset_folder, phase, target), IMAGES))
else:
logging.warning(f"Database for {
target} not found")
annotations_all["exists"] = annotations_all["name"].isin(exists_images)
annotations = annotations_all[annotations_all["exists"]]
return annotations
def run_convert(args: argparse.Namespace) -> None:
"""
Run convert
Parameters
----------
args: argparse.Namespace
arguments
Returns
-------
None
"""
conf = OmegaConf.load(args.cfg)
bbox_format = args.bbox_format
labels = {
label: num for (label, num) in zip(conf.dataset.targets, range(len(conf.dataset.targets)))}
dataset_folder = conf.dataset.dataset_folder
phases = conf.dataset.phases
for phase in phases:
phase_dir = os.path.join(dataset_folder, f"{
phase}_labels")
if not os.path.exists(phase_dir):
os.mkdir(phase_dir)
logging.info(f"Run convert {
phase}")
logging.info("Create Dataframe")
annotations = get_dataframe(conf, phase)
logging.info("Create image_path")
annotations["image_path"] = annotations.progress_apply(
lambda row: os.path.join(dataset_folder, phase, row["target"], row["name"]), axis=1
)
logging.info("Create keypoints")
annotations["keypoints"] = annotations.progress_apply(
lambda row: row["landmarks"], axis=1
)
if bbox_format == "cxcywh":
logging.info("Create bboxes cxcywh format")
annotations["converted_bboxes"] = annotations.progress_apply(
lambda row: xywh_to_cxcywh(row["bboxes"]), axis=1
)
elif bbox_format == "xyxy":
logging.info("Create bboxes xyxy format")
annotations["converted_bboxes"] = annotations.progress_apply(
lambda row: xywh_to_xyxy(row["bboxes"]), axis=1
)
elif bbox_format == "xywh":
logging.info("Create bboxes xywh format")
annotations["converted_bboxes"] = annotations["bboxes"]
else:
assert False, f"Unknown bbox format {
bbox_format}"
logging.info("Convert")
for target in annotations["target"].unique():
label_dir = os.path.join(phase_dir, target)
if not os.path.exists(label_dir):
os.mkdir(label_dir)
temp_annotations = annotations[annotations["target"] == target]
for row_id in tqdm(range(len(temp_annotations)), desc=target):
row = temp_annotations.iloc[row_id]
label_path = row["image_path"].replace(phase, f"{
phase}_labels").replace(".jpg", ".txt")
with open(label_path, "w") as f:
for i in range(len(row["labels"])):
if len(row["keypoints"][i]) ==0:
continue
f.write(str(0) + " ")
f.write(" ".join(map(str, row["converted_bboxes"][i])) +" ")
string =" ".join(map(str, row["keypoints"][i])).replace('[',"").replace(']',"").replace(',','')+"\n"
f.write(string)
if os.path.getsize(label_path) ==0:
os.remove(label_path)
os.remove(row["image_path"])
with open(f"{
phase}.txt", "w") as f:
for img_path in annotations["image_path"]:
f.write(f"{
img_path}\n")
if __name__ == "__main__":
parser = argparse.ArgumentParser("Convert Hagrid annotations to Yolo annotations format", add_help=False)
parser.add_argument("--bbox_format", default="cxcywh", type=str, help="bbox format: xyxy, cxcywh, xywh")
parser.add_argument("--cfg", default="converter_config.yaml", type=str, help="path to data config")
args = parser.parse_args()
run_convert(args)
#converter_config.yaml
dataset:
dataset_annotations: /home/lx/model_deploy/hagrid-Hagrid_v1/ # Path to annotations 修改成自己的
dataset_folder: /home/lx/model_deploy/hagrid-Hagrid_v1/ # Path to images 修改成自己的
phases: [train_val] #names of annotation directories 改成这个
targets:
- call
- dislike
- fist
- four
- like
- mute
- ok
- one
- palm
- peace
- rock
- stop
- stop_inverted
- three
- two_up
- two_up_inverted
- three2
- peace_inverted
- no_gesture
我对代码做了一定修改，原始数据集中有的标注中只有bbox没有keypoints，对这部分数据进行了清洗，执行以下命令转换成yolo数据集。
python -m converters.hagrid_to_yolo --cfg<PATH>
转换完成后会生成如下目录的图像和标注，然后我们将图片和数据集分别从各个类别中拿出来，生成一个包含所有类别图像的文件夹和包含所有标注的文件夹（也可以对上面的脚本直接改动，会方便很多，我为了验证数据集是否有问题为才继续按照这种分类划分的方式存储）。
将图像和标签的文件夹制作好以后，我们执行以下程序进行数据集划分，按照85%训练集，15%测试集的比例划分，自己可以更改。我们在划分时还额外进行一步判断，源图像和源标注是否存在，因为实际拆分时可能会发生有图像没有匹配标签的错误。
import os
import random
from shutil import copyfile
def split_dataset(images_dir, labels_dir, target_dir, train_ratio=0.85, val_ratio=0.15, seed=None):
"""
将数据集按照指定比例划分为训练集和验证集。
参数:
- images_dir: 包含图像文件的目录路径
- labels_dir: 包含标签文件的目录路径
- target_dir: 目标数据集文件夹路径，将包含划分后的 train 和 val 子目录
- train_ratio: 训练集比例
- val_ratio: 验证集比例
- seed: 随机种子，可选参数，用于确保随机性可复现
返回:
无返回值，划分后的数据集将保存在 target_dir 中
"""
# 获取图像文件列表
image_files = [file for file in os.listdir(images_dir) if file.endswith('.jpg')]
# 随机打乱文件顺序
random.seed(seed)
random.shuffle(image_files)
# 计算划分数据集的索引
total_files = len(image_files)
train_split = int(train_ratio * total_files)
val_split = int(val_ratio * total_files)
# 创建目标文件夹结构
os.makedirs(os.path.join(target_dir, 'train', 'images'), exist_ok=True)
os.makedirs(os.path.join(target_dir, 'train', 'labels'), exist_ok=True)
os.makedirs(os.path.join(target_dir, 'val', 'images'), exist_ok=True)
os.makedirs(os.path.join(target_dir, 'val', 'labels'), exist_ok=True)
# 划分数据集并复制到目标文件夹
for file in image_files[:train_split]:
src_image = os.path.join(images_dir, file)
src_label = os.path.join(labels_dir, file.replace('.jpg', '.txt'))
dst_image = os.path.join(target_dir, 'train', 'images', file)
dst_label = os.path.join(target_dir, 'train', 'labels', file.replace('.jpg', '.txt'))
if os.path.exists(scr_image):
copyfile(src_image, dst_image)
if os.path.exists(src_label):
copyfile(src_label, dst_label)
for file in image_files[train_split:train_split + val_split]:
src_image = os.path.join(images_dir, file)
src_label = os.path.join(labels_dir, file.replace('.jpg', '.txt'))
dst_image = os.path.join(target_dir, 'val', 'images', file)
dst_label = os.path.join(target_dir, 'val', 'labels', file.replace('.jpg', '.txt'))
if os.path.exists(scr_image):
copyfile(src_image, dst_image)
if os.path.exists(src_label):
copyfile(src_label, dst_label)
# 调用函数进行数据集划分
split_dataset(
images_dir='custom_images', #改为保存图像的文件夹
labels_dir='train_val_labels', #改为保存标注的文件夹
target_dir='dataset', #生成的数据集目标
train_ratio=0.85,
val_ratio=0.15,
seed=42 # 可选，设置随机种子以确保可复现性
)
然后我们得到了形如如下树形目录的yolo数据集：
dataset/
train/
images/
1.jpg
...
labels/
1.txt
...
val/
images/
1.jpg
...
labels/
1.txt
...
然后我们就可以开始训练啦^ _ ^