wehub-resource-sync
347 lines
12 KiB
Python
Executable File
347 lines
12 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Use the MediaPipe Gesture detection and Gesture landmark detection solutions to track hands and recognize gestures in images and videos."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import itertools
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Final
|
|
|
|
import cv2
|
|
import mediapipe as mp
|
|
import numpy as np
|
|
import requests
|
|
import tqdm
|
|
from mediapipe.tasks import python
|
|
from mediapipe.tasks.python import vision
|
|
|
|
import rerun as rr # pip install rerun-sdk
|
|
import rerun.blueprint as rrb
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterable
|
|
|
|
from mediapipe.tasks.python.components.containers import NormalizedLandmark
|
|
|
|
EXAMPLE_DIR: Final = Path(os.path.dirname(__file__))
|
|
DATASET_DIR: Final = EXAMPLE_DIR / "dataset" / "hand_gestures"
|
|
|
|
SAMPLE_IMAGE_PATH = EXAMPLE_DIR / "dataset" / "hand_gestures" / "victory.jpg"
|
|
|
|
# More samples: 'thumbs_down.jpg', 'victory.jpg', 'pointing_up.jpg', 'thumbs_up.jpg'
|
|
SAMPLE_IMAGE_URL = "https://storage.googleapis.com/mediapipe-tasks/gesture_recognizer/victory.jpg"
|
|
|
|
SAMPLE_VIDEO_PATH = EXAMPLE_DIR / "dataset" / "hand_gestures" / "peace.mp4"
|
|
|
|
SAMPLE_VIDEO_URL = "https://storage.googleapis.com/rerun-example-datasets/hand_gestures/peace.mp4"
|
|
|
|
# Emojis from https://github.com/googlefonts/noto-emoji/tree/main
|
|
GESTURE_URL = (
|
|
"https://raw.githubusercontent.com/googlefonts/noto-emoji/9cde38ef5ee6f090ce23f9035e494cb390a2b051/png/128/"
|
|
)
|
|
# Mapping of gesture categories to corresponding emojis
|
|
GESTURE_PICTURES = {
|
|
"None": "emoji_u2754.png",
|
|
"Closed_Fist": "emoji_u270a.png",
|
|
"Open_Palm": "emoji_u270b.png",
|
|
"Pointing_Up": "emoji_u261d.png",
|
|
"Thumb_Down": "emoji_u1f44e.png",
|
|
"Thumb_Up": "emoji_u1f44d.png",
|
|
"Victory": "emoji_u270c.png",
|
|
"ILoveYou": "emoji_u1f91f.png",
|
|
}
|
|
|
|
|
|
class GestureDetectorLogger:
|
|
"""
|
|
Logger for the MediaPipe Gesture Detection solution.
|
|
|
|
This class provides logging and utility functions for handling gesture recognition.
|
|
|
|
For more information on MediaPipe Gesture Detection:
|
|
<https://developers.google.com/mediapipe/solutions/vision/gesture_recognizer>
|
|
"""
|
|
|
|
# URL to the pre-trained MediaPipe Gesture Detection model
|
|
MODEL_DIR: Final = EXAMPLE_DIR / "model"
|
|
MODEL_PATH: Final = (MODEL_DIR / "gesture_recognizer.task").resolve()
|
|
MODEL_URL: Final = "https://storage.googleapis.com/mediapipe-models/gesture_recognizer/gesture_recognizer/float16/latest/gesture_recognizer.task"
|
|
|
|
def __init__(self, video_mode: bool = False) -> None:
|
|
self._video_mode = video_mode
|
|
|
|
if not self.MODEL_PATH.exists():
|
|
download_file(self.MODEL_URL, self.MODEL_PATH)
|
|
|
|
base_options = python.BaseOptions(model_asset_path=str(self.MODEL_PATH))
|
|
options = vision.GestureRecognizerOptions(
|
|
base_options=base_options,
|
|
running_mode=mp.tasks.vision.RunningMode.VIDEO if self._video_mode else mp.tasks.vision.RunningMode.IMAGE,
|
|
)
|
|
self.recognizer = vision.GestureRecognizer.create_from_options(options)
|
|
|
|
rr.log(
|
|
"/",
|
|
rr.AnnotationContext(
|
|
rr.ClassDescription(
|
|
info=rr.AnnotationInfo(id=0, label="Hand3D"),
|
|
keypoint_connections=mp.solutions.hands.HAND_CONNECTIONS,
|
|
),
|
|
),
|
|
static=True,
|
|
)
|
|
rr.log("hand3d", rr.ViewCoordinates.LEFT_HAND_Y_DOWN, static=True)
|
|
|
|
@staticmethod
|
|
def convert_landmarks_to_image_coordinates(
|
|
hand_landmarks: list[list[NormalizedLandmark]],
|
|
width: int,
|
|
height: int,
|
|
) -> list[tuple[int, int]]:
|
|
return [(int(lm.x * width), int(lm.y * height)) for hand_landmark in hand_landmarks for lm in hand_landmark]
|
|
|
|
@staticmethod
|
|
def convert_landmarks_to_3d(hand_landmarks: list[list[NormalizedLandmark]]) -> list[tuple[float, float, float]]:
|
|
return [(lm.x, lm.y, lm.z) for hand_landmark in hand_landmarks for lm in hand_landmark]
|
|
|
|
def detect_and_log(self, image: cv2.typing.MatLike, frame_time_nano: int) -> None:
|
|
# Recognize gestures in the image
|
|
height, width, _ = image.shape
|
|
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
|
|
|
|
recognition_result = (
|
|
self.recognizer.recognize_for_video(image, int(frame_time_nano / 1e6))
|
|
if self._video_mode
|
|
else self.recognizer.recognize(image)
|
|
)
|
|
|
|
for log_key in ["hand2d/points", "hand2d/connections", "hand3d/points"]:
|
|
rr.log(log_key, rr.Clear(recursive=True))
|
|
|
|
for gesture in recognition_result.gestures:
|
|
# Get the top gesture from the recognition result
|
|
gesture_category = gesture[0].category_name if recognition_result.gestures else "None"
|
|
self.present_detected_gesture(gesture_category) # Log the detected gesture
|
|
|
|
if recognition_result.hand_landmarks:
|
|
hand_landmarks = recognition_result.hand_landmarks
|
|
|
|
landmark_positions_3d = self.convert_landmarks_to_3d(hand_landmarks)
|
|
if landmark_positions_3d is not None:
|
|
rr.log(
|
|
"hand3d/points",
|
|
rr.Points3D(
|
|
landmark_positions_3d,
|
|
radii=20,
|
|
class_ids=0,
|
|
keypoint_ids=list(range(len(landmark_positions_3d))),
|
|
),
|
|
)
|
|
|
|
# Convert normalized coordinates to image coordinates
|
|
points = self.convert_landmarks_to_image_coordinates(hand_landmarks, width, height)
|
|
|
|
# Log points to the image and Hand Entity
|
|
rr.log("hand2d/points", rr.Points2D(points, radii=10, colors=[255, 0, 0]))
|
|
|
|
# Obtain hand connections from MediaPipe
|
|
mp_hands_connections = mp.solutions.hands.HAND_CONNECTIONS
|
|
points1 = [points[connection[0]] for connection in mp_hands_connections]
|
|
points2 = [points[connection[1]] for connection in mp_hands_connections]
|
|
|
|
# Log connections to the image and Hand Entity [128, 128, 128]
|
|
rr.log("hand2d/connections", rr.LineStrips2D(np.stack((points1, points2), axis=1), colors=[255, 165, 0]))
|
|
|
|
def present_detected_gesture(self, category: str) -> None:
|
|
# Get the corresponding ulr of the picture for the detected gesture category
|
|
gesture_pic = GESTURE_PICTURES.get(
|
|
category,
|
|
"emoji_u2754.png", # default
|
|
)
|
|
|
|
# Log the detection by using the appropriate image
|
|
rr.log(
|
|
"detection",
|
|
rr.TextDocument(f"".strip(), media_type=rr.MediaType.MARKDOWN),
|
|
)
|
|
|
|
|
|
def download_file(url: str, path: Path) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
logging.info("Downloading %s to %s", url, path)
|
|
response = requests.get(url, stream=True)
|
|
with tqdm.tqdm.wrapattr(
|
|
open(path, "wb"),
|
|
"write",
|
|
miniters=1,
|
|
total=int(response.headers.get("content-length", 0)),
|
|
desc=f"Downloading {path.name}",
|
|
) as f:
|
|
for chunk in response.iter_content(chunk_size=4096):
|
|
f.write(chunk)
|
|
|
|
|
|
def resize_image(image: cv2.typing.MatLike, max_dim: int | None) -> cv2.typing.MatLike:
|
|
"""Resize an image if it is larger than max_dim."""
|
|
if max_dim is None:
|
|
return image
|
|
height, width, _ = image.shape
|
|
scale = max_dim / max(height, width)
|
|
if scale < 1:
|
|
image = cv2.resize(image, (0, 0), fx=scale, fy=scale)
|
|
return image
|
|
|
|
|
|
def run_from_sample_image(path: Path | str) -> None:
|
|
"""Run the gesture recognition on a single image."""
|
|
image = cv2.imread(str(path))
|
|
# image = resize_image(image, max_dim)
|
|
rr.log("media/image", rr.Image(image, color_model="BGR"))
|
|
|
|
detect_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
logger = GestureDetectorLogger(video_mode=False)
|
|
logger.detect_and_log(detect_image, 0)
|
|
|
|
|
|
def run_from_video_capture(vid: int | str, max_frame_count: int | None) -> None:
|
|
"""
|
|
Run the detector on a video stream.
|
|
|
|
Parameters
|
|
----------
|
|
vid:
|
|
The video stream to run the detector on. Use 0/1 for the default camera or a path to a video file.
|
|
max_frame_count:
|
|
The maximum number of frames to process. If None, process all frames.
|
|
|
|
"""
|
|
cap = cv2.VideoCapture(vid)
|
|
fps = cap.get(cv2.CAP_PROP_FPS)
|
|
|
|
detector = GestureDetectorLogger(video_mode=True)
|
|
|
|
try:
|
|
it: Iterable[int] = itertools.count() if max_frame_count is None else range(max_frame_count)
|
|
|
|
for frame_idx in tqdm.tqdm(it, desc="Processing frames"):
|
|
# Capture frame-by-frame
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# OpenCV sometimes returns a blank frame, so we skip it
|
|
if np.all(frame == 0):
|
|
continue
|
|
|
|
# frame = resize_image(frame, max_dim)
|
|
|
|
# get frame time
|
|
frame_time_nano = int(cap.get(cv2.CAP_PROP_POS_MSEC) * 1e6)
|
|
if frame_time_nano == 0:
|
|
# On some platforms it always returns zero, so we compute from the frame counter and fps
|
|
frame_time_nano = int(frame_idx * 1000 / fps * 1e6)
|
|
|
|
# log data
|
|
rr.set_time("frame_nr", sequence=frame_idx)
|
|
rr.set_time("frame_time", duration=1e-9 * frame_time_nano)
|
|
detector.detect_and_log(frame, frame_time_nano)
|
|
rr.log("media/video", rr.Image(frame, color_model="BGR").compress(jpeg_quality=75))
|
|
|
|
except KeyboardInterrupt:
|
|
pass
|
|
|
|
# When everything done, release the capture
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
def main() -> None:
|
|
# Ensure the logging gets written to stderr
|
|
logging.getLogger().addHandler(logging.StreamHandler())
|
|
logging.getLogger().setLevel("INFO")
|
|
|
|
# Set up argument parser with description
|
|
parser = argparse.ArgumentParser(
|
|
description="Uses the MediaPipe Gesture Recognition to track a hand and recognize gestures in image or video.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--demo-image",
|
|
action="store_true",
|
|
help="Run on a demo image automatically downloaded",
|
|
)
|
|
parser.add_argument("--demo-video", action="store_true", help="Run on a demo image automatically downloaded.")
|
|
parser.add_argument(
|
|
"--image",
|
|
type=Path,
|
|
help="Run on the provided image",
|
|
)
|
|
parser.add_argument("--video", type=Path, help="Run on the provided video file.")
|
|
parser.add_argument(
|
|
"--camera",
|
|
type=int,
|
|
default=0,
|
|
help="Run from the camera stream (parameter is the camera ID, usually 0; or maybe 1 on mac)",
|
|
)
|
|
parser.add_argument(
|
|
"--max-frame",
|
|
type=int,
|
|
help="Stop after processing this many frames. If not specified, will run until interrupted.",
|
|
)
|
|
|
|
# Add Rerun specific arguments
|
|
rr.script_add_args(parser)
|
|
|
|
# Parse command line arguments
|
|
args, unknown = parser.parse_known_args()
|
|
for arg in unknown: # Log any unknown arguments
|
|
logging.warning(f"unknown arg: {arg}")
|
|
|
|
# Set up Rerun with script name
|
|
rr.script_setup(
|
|
args,
|
|
"rerun_example_mp_gesture_recognition",
|
|
default_blueprint=rrb.Horizontal(
|
|
rrb.Spatial2DView(name="Input & Hand", contents=["media/**", "hand2d/**"]),
|
|
rrb.Vertical(
|
|
rrb.Tabs(
|
|
rrb.Spatial3DView(name="Hand 3D", origin="hand3d"),
|
|
rrb.Spatial2DView(name="Hand 2D", origin="hand2d"),
|
|
),
|
|
rrb.TextDocumentView(name="Detection", origin="detection"),
|
|
row_shares=[3, 2],
|
|
),
|
|
column_shares=[3, 1],
|
|
),
|
|
)
|
|
|
|
# Choose the appropriate run mode based on provided arguments
|
|
if args.demo_image:
|
|
if not SAMPLE_IMAGE_PATH.exists():
|
|
download_file(SAMPLE_IMAGE_URL, SAMPLE_IMAGE_PATH)
|
|
run_from_sample_image(SAMPLE_IMAGE_PATH)
|
|
elif args.demo_video:
|
|
if not SAMPLE_VIDEO_PATH.exists():
|
|
download_file(SAMPLE_VIDEO_URL, SAMPLE_VIDEO_PATH)
|
|
run_from_video_capture(str(SAMPLE_VIDEO_PATH), args.max_frame)
|
|
elif args.image:
|
|
run_from_sample_image(args.image)
|
|
elif args.video:
|
|
run_from_video_capture(args.video, args.max_frame)
|
|
elif args.camera:
|
|
run_from_video_capture(int(args.camera), args.max_frame)
|
|
else:
|
|
if not SAMPLE_VIDEO_PATH.exists():
|
|
download_file(SAMPLE_VIDEO_URL, SAMPLE_VIDEO_PATH)
|
|
run_from_video_capture(str(SAMPLE_VIDEO_PATH), args.max_frame)
|
|
|
|
# Tear down Rerun script
|
|
rr.script_teardown(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|