chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+518
@@ -0,0 +1,518 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Use the MediaPipe Face detection and Face landmark detection solutions to track human faces in images and videos."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import itertools
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
import numpy as np
|
||||
import requests
|
||||
import tqdm
|
||||
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, Iterator
|
||||
|
||||
# If set, log everything as static.
|
||||
#
|
||||
# Generally, the Viewer accumulates data until its set memory budget at which point it will
|
||||
# remove the oldest data from the recording (see https://rerun.io/docs/howto/visualization/limit-ram)
|
||||
# By instead logging data as static, no data will be accumulated over time since previous
|
||||
# data is overwritten.
|
||||
# Naturally, the drawback of this is that there's no history of previous data sent to the viewer,
|
||||
# as well as no timestamps, making the Viewer's timeline effectively inactive.
|
||||
ALL_STATIC: bool = False
|
||||
|
||||
EXAMPLE_DIR: Final = Path(os.path.dirname(__file__))
|
||||
DATASET_DIR: Final = EXAMPLE_DIR / "dataset"
|
||||
MODEL_DIR: Final = EXAMPLE_DIR / "model"
|
||||
|
||||
SAMPLE_IMAGE_PATH = (DATASET_DIR / "image.jpg").resolve()
|
||||
# from https://pixabay.com/photos/brother-sister-girl-family-boy-977170/
|
||||
SAMPLE_IMAGE_URL = "https://i.imgur.com/Vu2Nqwb.jpg"
|
||||
|
||||
# uncomment blendshapes of interest
|
||||
BLENDSHAPES_CATEGORIES = {
|
||||
"_neutral",
|
||||
"browDownLeft",
|
||||
"browDownRight",
|
||||
"browInnerUp",
|
||||
"browOuterUpLeft",
|
||||
"browOuterUpRight",
|
||||
"cheekPuff",
|
||||
"cheekSquintLeft",
|
||||
"cheekSquintRight",
|
||||
"eyeBlinkLeft",
|
||||
"eyeBlinkRight",
|
||||
"eyeLookDownLeft",
|
||||
"eyeLookDownRight",
|
||||
"eyeLookInLeft",
|
||||
"eyeLookInRight",
|
||||
"eyeLookOutLeft",
|
||||
"eyeLookOutRight",
|
||||
"eyeLookUpLeft",
|
||||
"eyeLookUpRight",
|
||||
"eyeSquintLeft",
|
||||
"eyeSquintRight",
|
||||
"eyeWideLeft",
|
||||
"eyeWideRight",
|
||||
"jawForward",
|
||||
"jawLeft",
|
||||
"jawOpen",
|
||||
"jawRight",
|
||||
"mouthClose",
|
||||
"mouthDimpleLeft",
|
||||
"mouthDimpleRight",
|
||||
"mouthFrownLeft",
|
||||
"mouthFrownRight",
|
||||
"mouthFunnel",
|
||||
"mouthLeft",
|
||||
"mouthLowerDownLeft",
|
||||
"mouthLowerDownRight",
|
||||
"mouthPressLeft",
|
||||
"mouthPressRight",
|
||||
"mouthPucker",
|
||||
"mouthRight",
|
||||
"mouthRollLower",
|
||||
"mouthRollUpper",
|
||||
"mouthShrugLower",
|
||||
"mouthShrugUpper",
|
||||
"mouthSmileLeft",
|
||||
"mouthSmileRight",
|
||||
"mouthStretchLeft",
|
||||
"mouthStretchRight",
|
||||
"mouthUpperUpLeft",
|
||||
"mouthUpperUpRight",
|
||||
"noseSneerLeft",
|
||||
"noseSneerRight",
|
||||
}
|
||||
|
||||
|
||||
class FaceDetectorLogger:
|
||||
"""
|
||||
Logger for the MediaPipe Face Detection solution.
|
||||
|
||||
<https://developers.google.com/mediapipe/solutions/vision/face_detector>
|
||||
"""
|
||||
|
||||
MODEL_PATH: Final = (MODEL_DIR / "blaze_face_short_range.tflite").resolve()
|
||||
MODEL_URL: Final = (
|
||||
"https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/"
|
||||
"blaze_face_short_range.tflite"
|
||||
)
|
||||
|
||||
def __init__(self, video_mode: bool = False) -> None:
|
||||
self._video_mode = video_mode
|
||||
|
||||
# download model if necessary
|
||||
if not self.MODEL_PATH.exists():
|
||||
download_file(self.MODEL_URL, self.MODEL_PATH)
|
||||
|
||||
self._base_options = mp.tasks.BaseOptions(
|
||||
model_asset_path=str(self.MODEL_PATH),
|
||||
)
|
||||
self._options = vision.FaceDetectorOptions(
|
||||
base_options=self._base_options,
|
||||
running_mode=mp.tasks.vision.RunningMode.VIDEO if self._video_mode else mp.tasks.vision.RunningMode.IMAGE,
|
||||
)
|
||||
self._detector = vision.FaceDetector.create_from_options(self._options)
|
||||
|
||||
# With this annotation, the viewer will connect the keypoints with some lines to improve visibility.
|
||||
rr.log(
|
||||
"video/detector",
|
||||
rr.ClassDescription(
|
||||
info=rr.AnnotationInfo(id=0),
|
||||
keypoint_connections=[(0, 1), (1, 2), (2, 0), (2, 3), (0, 4), (1, 5)],
|
||||
),
|
||||
static=True,
|
||||
)
|
||||
|
||||
def detect_and_log(self, image: cv2.typing.MatLike, frame_time_nano: int) -> None:
|
||||
height, width, _ = image.shape
|
||||
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
|
||||
|
||||
detection_result = (
|
||||
self._detector.detect_for_video(image, int(frame_time_nano / 1e6))
|
||||
if self._video_mode
|
||||
else self._detector.detect(image)
|
||||
)
|
||||
rr.log("video/detector/faces", rr.Clear(recursive=True), static=ALL_STATIC)
|
||||
for i, detection in enumerate(detection_result.detections):
|
||||
# log bounding box
|
||||
bbox = detection.bounding_box
|
||||
index, score = detection.categories[0].index, detection.categories[0].score
|
||||
|
||||
# log bounding box
|
||||
rr.log(
|
||||
f"video/detector/faces/{i}/bbox",
|
||||
rr.Boxes2D(
|
||||
array=[bbox.origin_x, bbox.origin_y, bbox.width, bbox.height],
|
||||
array_format=rr.Box2DFormat.XYWH,
|
||||
),
|
||||
rr.AnyValues(index=index, score=score),
|
||||
static=ALL_STATIC,
|
||||
)
|
||||
|
||||
# MediaPipe's keypoints are normalized to [0, 1], so we need to scale them to get pixel coordinates.
|
||||
pts = [
|
||||
(math.floor(keypoint.x * width), math.floor(keypoint.y * height)) for keypoint in detection.keypoints
|
||||
]
|
||||
rr.log(
|
||||
f"video/detector/faces/{i}/keypoints",
|
||||
rr.Points2D(pts, radii=3, keypoint_ids=list(range(6))),
|
||||
static=ALL_STATIC,
|
||||
)
|
||||
|
||||
|
||||
class FaceLandmarkerLogger:
|
||||
"""
|
||||
Logger for the MediaPipe Face Landmark Detection solution.
|
||||
|
||||
<https://developers.google.com/mediapipe/solutions/vision/face_landmarker>
|
||||
"""
|
||||
|
||||
MODEL_PATH: Final = (MODEL_DIR / "face_landmarker.task").resolve()
|
||||
MODEL_URL: Final = (
|
||||
"https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/"
|
||||
"face_landmarker.task"
|
||||
)
|
||||
|
||||
def __init__(self, video_mode: bool = False, num_faces: int = 1) -> None:
|
||||
self._video_mode = video_mode
|
||||
|
||||
# download model if necessary
|
||||
if not self.MODEL_PATH.exists():
|
||||
download_file(self.MODEL_URL, self.MODEL_PATH)
|
||||
|
||||
self._base_options = mp.tasks.BaseOptions(
|
||||
model_asset_path=str(self.MODEL_PATH),
|
||||
)
|
||||
self._options = vision.FaceLandmarkerOptions(
|
||||
base_options=self._base_options,
|
||||
output_face_blendshapes=True,
|
||||
num_faces=num_faces,
|
||||
running_mode=mp.tasks.vision.RunningMode.VIDEO if self._video_mode else mp.tasks.vision.RunningMode.IMAGE,
|
||||
)
|
||||
self._detector = vision.FaceLandmarker.create_from_options(self._options)
|
||||
|
||||
# Extract classes from MediaPipe face mesh solution. The goal of this code is:
|
||||
# 1) Log an annotation context with one class ID per facial feature. For each class ID, the class description
|
||||
# contains the connections between corresponding keypoints (taken from the MediaPipe face mesh solution)
|
||||
# 2) A class ID array matching the class IDs in the annotation context to keypoint indices (to be passed as
|
||||
# the `class_ids` argument to `rr.log`).
|
||||
|
||||
classes = [
|
||||
mp.solutions.face_mesh.FACEMESH_LIPS,
|
||||
mp.solutions.face_mesh.FACEMESH_LEFT_EYE,
|
||||
mp.solutions.face_mesh.FACEMESH_LEFT_IRIS,
|
||||
mp.solutions.face_mesh.FACEMESH_LEFT_EYEBROW,
|
||||
mp.solutions.face_mesh.FACEMESH_RIGHT_EYE,
|
||||
mp.solutions.face_mesh.FACEMESH_RIGHT_EYEBROW,
|
||||
mp.solutions.face_mesh.FACEMESH_RIGHT_IRIS,
|
||||
mp.solutions.face_mesh.FACEMESH_FACE_OVAL,
|
||||
mp.solutions.face_mesh.FACEMESH_NOSE,
|
||||
]
|
||||
|
||||
self._class_ids = [0] * mp.solutions.face_mesh.FACEMESH_NUM_LANDMARKS_WITH_IRISES
|
||||
class_descriptions = []
|
||||
for i, klass in enumerate(classes):
|
||||
# MediaPipe only provides connections for class, not actual class per keypoint. So we have to extract the
|
||||
# classes from the connections.
|
||||
ids = set()
|
||||
for connection in klass:
|
||||
ids.add(connection[0])
|
||||
ids.add(connection[1])
|
||||
|
||||
for id_ in ids:
|
||||
self._class_ids[id_] = i
|
||||
|
||||
class_descriptions.append(
|
||||
rr.ClassDescription(
|
||||
info=rr.AnnotationInfo(id=i),
|
||||
keypoint_connections=klass,
|
||||
),
|
||||
)
|
||||
|
||||
rr.log("video/landmarker", rr.AnnotationContext(class_descriptions), static=True)
|
||||
rr.log("reconstruction", rr.AnnotationContext(class_descriptions), static=True)
|
||||
|
||||
# properly align the 3D face in the viewer
|
||||
rr.log("reconstruction", rr.ViewCoordinates.RDF, static=True)
|
||||
|
||||
def detect_and_log(self, image: cv2.typing.MatLike, frame_time_nano: int) -> None:
|
||||
height, width, _ = image.shape
|
||||
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
|
||||
|
||||
detection_result = (
|
||||
self._detector.detect_for_video(image, int(frame_time_nano / 1e6))
|
||||
if self._video_mode
|
||||
else self._detector.detect(image)
|
||||
)
|
||||
|
||||
def is_empty(i: Iterator[Any]) -> bool:
|
||||
try:
|
||||
next(i)
|
||||
return False
|
||||
except StopIteration:
|
||||
return True
|
||||
|
||||
if is_empty(zip(detection_result.face_landmarks, detection_result.face_blendshapes, strict=False)):
|
||||
rr.log("video/landmarker/faces", rr.Clear(recursive=True), static=ALL_STATIC)
|
||||
rr.log("reconstruction/faces", rr.Clear(recursive=True), static=ALL_STATIC)
|
||||
rr.log("blendshapes", rr.Clear(recursive=True), static=ALL_STATIC)
|
||||
|
||||
for i, (landmark, blendshapes) in enumerate(
|
||||
zip(detection_result.face_landmarks, detection_result.face_blendshapes, strict=False),
|
||||
):
|
||||
if len(landmark) == 0 or len(blendshapes) == 0:
|
||||
rr.log(
|
||||
f"video/landmarker/faces/{i}/landmarks",
|
||||
rr.Clear(recursive=True),
|
||||
static=ALL_STATIC,
|
||||
)
|
||||
rr.log(
|
||||
f"reconstruction/faces/{i}",
|
||||
rr.Clear(recursive=True),
|
||||
static=ALL_STATIC,
|
||||
)
|
||||
rr.log(f"blendshapes/{i}", rr.Clear(recursive=True), static=ALL_STATIC)
|
||||
continue
|
||||
|
||||
# MediaPipe's keypoints are normalized to [0, 1], so we need to scale them to get pixel coordinates.
|
||||
pts = [(math.floor(lm.x * width), math.floor(lm.y * height)) for lm in landmark]
|
||||
keypoint_ids = list(range(len(landmark)))
|
||||
rr.log(
|
||||
f"video/landmarker/faces/{i}/landmarks",
|
||||
rr.Points2D(pts, radii=3, keypoint_ids=keypoint_ids, class_ids=self._class_ids),
|
||||
static=ALL_STATIC,
|
||||
)
|
||||
|
||||
rr.log(
|
||||
f"reconstruction/faces/{i}",
|
||||
rr.Points3D(
|
||||
[(lm.x, lm.y, lm.z) for lm in landmark],
|
||||
keypoint_ids=keypoint_ids,
|
||||
class_ids=self._class_ids,
|
||||
),
|
||||
static=ALL_STATIC,
|
||||
)
|
||||
|
||||
for blendshape in blendshapes:
|
||||
if blendshape.category_name in BLENDSHAPES_CATEGORIES:
|
||||
# NOTE(cmc): That one we still log as temporal, otherwise it's really meh.
|
||||
rr.log(
|
||||
f"blendshapes/{i}/{blendshape.category_name}",
|
||||
rr.Scalars(blendshape.score),
|
||||
)
|
||||
|
||||
|
||||
# ========================================================================================
|
||||
# Main & CLI code
|
||||
|
||||
|
||||
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_video_capture(vid: int | str, max_dim: int | None, max_frame_count: int | None, num_faces: int) -> None:
|
||||
"""
|
||||
Run the face detector on a video stream.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vid:
|
||||
The video stream to run the detector on. Use 0 for the default camera or a path to a video file.
|
||||
max_dim:
|
||||
The maximum dimension of the image. If the image is larger, it will be scaled down.
|
||||
max_frame_count:
|
||||
The maximum number of frames to process. If None, process all frames.
|
||||
num_faces:
|
||||
The number of faces to track. If set to 1, temporal smoothing will be applied.
|
||||
|
||||
"""
|
||||
|
||||
cap = cv2.VideoCapture(vid)
|
||||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||||
|
||||
detector = FaceDetectorLogger(video_mode=True)
|
||||
landmarker = FaceLandmarkerLogger(video_mode=True, num_faces=num_faces)
|
||||
|
||||
print("Capturing video stream. Press ctrl-c to stop.")
|
||||
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)
|
||||
landmarker.detect_and_log(frame, frame_time_nano)
|
||||
rr.log(
|
||||
"video/image",
|
||||
rr.Image(frame, color_model="BGR"),
|
||||
static=ALL_STATIC,
|
||||
)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
# When everything done, release the capture
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
def run_from_sample_image(path: Path, max_dim: int | None, num_faces: int) -> None:
|
||||
"""Run the face detector on a single image."""
|
||||
image = cv2.imread(str(path))
|
||||
image = resize_image(image, max_dim)
|
||||
logger = FaceDetectorLogger(video_mode=False)
|
||||
landmarker = FaceLandmarkerLogger(video_mode=False, num_faces=num_faces)
|
||||
logger.detect_and_log(image, 0)
|
||||
landmarker.detect_and_log(image, 0)
|
||||
rr.log(
|
||||
"video/image",
|
||||
rr.Image(image, color_model="BGR"),
|
||||
static=ALL_STATIC,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
logging.getLogger().addHandler(logging.StreamHandler())
|
||||
logging.getLogger().setLevel("INFO")
|
||||
|
||||
parser = argparse.ArgumentParser(description="Uses the MediaPipe Face Detection to track a human pose in video.")
|
||||
parser.add_argument(
|
||||
"--demo-image",
|
||||
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)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-frame",
|
||||
type=int,
|
||||
help="Stop after processing this many frames. If not specified, will run until interrupted.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-dim",
|
||||
type=int,
|
||||
help="Resize the image such as its maximum dimension is not larger than this value.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-faces",
|
||||
type=int,
|
||||
default=1,
|
||||
help=(
|
||||
"Max number of faces detected by the landmark model (temporal smoothing is applied only for a value of 1)."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--static", action="store_true", help="If set, logs everything as static")
|
||||
|
||||
rr.script_add_args(parser)
|
||||
|
||||
args, unknown = parser.parse_known_args()
|
||||
for arg in unknown:
|
||||
logging.warning(f"unknown arg: {arg}")
|
||||
|
||||
rr.script_setup(
|
||||
args,
|
||||
"rerun_example_mp_face_detection",
|
||||
default_blueprint=rrb.Horizontal(
|
||||
rrb.Spatial3DView(origin="reconstruction"),
|
||||
rrb.Vertical(
|
||||
rrb.Spatial2DView(origin="video"),
|
||||
rrb.TimeSeriesView(
|
||||
origin="blendshapes",
|
||||
# Enable only certain blend shapes by default. More can be added in the viewer ui
|
||||
contents=[
|
||||
"+ blendshapes/0/eyeBlinkLeft",
|
||||
"+ blendshapes/0/eyeBlinkRight",
|
||||
"+ blendshapes/0/jawOpen",
|
||||
"+ blendshapes/0/mouthSmileLeft",
|
||||
"+ blendshapes/0/mouthSmileRight",
|
||||
],
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
global ALL_STATIC
|
||||
ALL_STATIC = args.static
|
||||
|
||||
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, args.max_dim, args.num_faces)
|
||||
elif args.image is not None:
|
||||
run_from_sample_image(args.image, args.max_dim, args.num_faces)
|
||||
elif args.video is not None:
|
||||
run_from_video_capture(str(args.video), args.max_dim, args.max_frame, args.num_faces)
|
||||
else:
|
||||
run_from_video_capture(args.camera, args.max_dim, args.max_frame, args.num_faces)
|
||||
|
||||
rr.script_teardown(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user