MediaPipe自定义手势识别训练模型

前言¶

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.001
batch_size：批大小，默认为2
epochs：训练轮数，默认为10
shuffle：是否打乱数据，默认为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自定义手势识别训练模型】获取源码

前言¶