前言

MediaPipe是Google开源的一个用于构建感知管道以处理视频、音频等时间序列数据的框架。其中MediaPipe Hands是一个高性能的手部关键点检测解决方案,能够在移动设备上实时检测手部关键点。本项目主要基于MediaPipe Model Maker来训练自定义手势识别模型,支持石头、剪刀、布、无手势等多种手势的识别。项目提供了完整的训练和推理流程,支持图片和视频两种推理模式,可以实时检测和识别手势。

使用环境:

  • Python 3.11
  • TensorFlow 2.15.1
  • MediaPipe 0.10.21
  • Ubuntu 22.04

项目主要程序介绍

  1. train.py:训练自定义手势识别模型。
  2. infer.py:使用训练好的模型进行手势识别推理,支持图片和视频模式。
  3. utils.py:工具函数,包含手部关键点绘制和手势信息显示功能。
  4. models/:存放预训练模型和自定义训练模型的目录。
  5. 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自定义手势识别训练模型】获取源码

参考资料

  1. 自定义手势模型
  2. Python 手势识别指南
  3. TensorFlow
Logo

技术共进,成长同行——讯飞AI开发者社区

更多推荐