MediaPipe自定义手势识别训练模型
摘要:本项目基于MediaPipe Model Maker框架开发了一个自定义手势识别系统,支持石头剪刀布等多种手势的实时检测。系统包含完整的训练流程(支持调整学习率、批次大小等参数)和两种推理模式(图片/视频)。使用Python 3.11环境,通过100-200张手势图片训练后,模型准确率达95%。项目提供源码下载和详细文档,适用于移动设备上的实时手势交互应用。
前言
MediaPipe是Google开源的一个用于构建感知管道以处理视频、音频等时间序列数据的框架。其中MediaPipe Hands是一个高性能的手部关键点检测解决方案,能够在移动设备上实时检测手部关键点。本项目主要基于MediaPipe Model Maker来训练自定义手势识别模型,支持石头、剪刀、布、无手势等多种手势的识别。项目提供了完整的训练和推理流程,支持图片和视频两种推理模式,可以实时检测和识别手势。
使用环境:
- Python 3.11
- TensorFlow 2.15.1
- MediaPipe 0.10.21
- Ubuntu 22.04
项目主要程序介绍
train.py
:训练自定义手势识别模型。infer.py
:使用训练好的模型进行手势识别推理,支持图片和视频模式。utils.py
:工具函数,包含手部关键点绘制和手势信息显示功能。models/
:存放预训练模型和自定义训练模型的目录。dataset/
:训练数据集目录,包含不同手势的图片数据。
安装环境
-
首先安装Python环境,推荐使用Python 3.8或更高版本。
-
安装所需的依赖库。
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
准备数据
训练数据集存放在dataset/images/
目录下,按照不同的手势类型分别存放在不同的子目录中,自定义的数据集可以存放在这个位置,根据自己的清空创建更多的文件夹,图片是包含一个手势:
dataset/
- images/
- rock/ # 石头手势图片
- paper/ # 布手势图片
- scissors/ # 剪刀手势图片
- none/ # 无手势/背景图片
每个手势类型建议准备100-200张不同角度、不同光照条件下的图片,以提高模型的泛化能力。图片格式支持jpg、png等常见格式。
训练模型
准备好数据之后,就可以开始训练自定义手势识别模型了。训练程序会自动读取数据集,并使用MediaPipe Model Maker进行训练。
训练参数说明
在train.py
中可以调整以下主要参数:
learning_rate
:学习率,默认为0.001batch_size
:批大小,默认为2epochs
:训练轮数,默认为10shuffle
:是否打乱数据,默认为True
开始训练
训练的代码片段:
from mediapipe_model_maker import gesture_recognizer
# 读取数据集
data = gesture_recognizer.Dataset.from_folder(
dirname="dataset/images",
hparams=gesture_recognizer.HandDataPreprocessingParams()
)
# 数据分割
train_data, rest_data = data.split(0.8)
validation_data, test_data = rest_data.split(0.5)
# 设置训练参数
hparams = gesture_recognizer.HParams(
learning_rate=0.001,
batch_size=2,
epochs=10,
shuffle=True
)
# 训练模型
model = gesture_recognizer.GestureRecognizer.create(
train_data=train_data,
validation_data=validation_data,
options=gesture_recognizer.GestureRecognizerOptions(hparams=hparams)
)
# 评估和导出
loss, acc = model.evaluate(test_data, batch_size=1)
model.export_model(model_name="custom_gesture_recognizer.task")
完整的代码在train.py
。
python train.py
训练过程中会显示训练进度和损失值:
标签列表:['none', 'paper', 'rock', 'scissors']
Model checkpoint saved at exported_model/checkpoint
Training finished.
Test loss:0.123, Test accuracy:0.95
训练完成后,模型会自动保存到models/custom_gesture_recognizer.task
。
模型推理
项目支持两种推理模式:图片推理和视频推理。
图片推理
对单张图片进行手势识别,简单的代码片段。
import cv2
import mediapipe as mp
from utils import draw_landmarks_and_connections, draw_gesture_info, HAND_CONNECTIONS
# 创建手势识别器
BaseOptions = mp.tasks.BaseOptions
GestureRecognizer = mp.tasks.vision.GestureRecognizer
GestureRecognizerOptions = mp.tasks.vision.GestureRecognizerOptions
VisionRunningMode = mp.tasks.vision.RunningMode
options = GestureRecognizerOptions(
base_options=BaseOptions(model_asset_path='models/custom_gesture_recognizer.task'),
num_hands=2,
running_mode=VisionRunningMode.IMAGE)
with GestureRecognizer.create_from_options(options) as recognizer:
# 读取图片
mp_image = mp.Image.create_from_file('test.jpg')
result = recognizer.recognize(mp_image)
# 输出识别结果
for i in range(len(result.gestures)):
print(f"第{i + 1}个手势:{result.gestures[i][0].category_name},置信度:{result.gestures[i][0].score:.1%}")
完整代码在infer.py
python infer.py
默认配置为图片模式,会对test2.jpg
进行推理。可以在infer.py
中修改以下参数:
mode
:设置为"image"model_asset_path
:模型路径,可选择官方模型或自定义模型show_image_or_video
:是否显示结果图片
视频推理
对摄像头实时视频流进行手势识别:
# 修改infer.py中的mode为"video"
python infer.py
视频模式相关参数:
mode
:设置为"video"video_path_or_id
:视频文件路径或摄像头ID(0为默认摄像头)show_image_or_video
:是否显示实时视频窗口
推理输出
推理程序会输出以下信息:
第1个手势:rock,置信度:95.2%
第2个手势:scissors,置信度:87.1%
结果图像已保存为 result_with_landmarks.jpg
同时会在图片上绘制手部关键点和手势识别结果。
获取源码
在微信公众号回复【MediaPipe自定义手势识别训练模型】获取源码
参考资料
更多推荐
所有评论(0)