"""Drive a rigged 3D character with body pose keypoints from a camera.""" from __future__ import annotations import argparse import sys import time from pathlib import Path from typing import NamedTuple import cv2 import numpy as np import mesh3d import ui from camera import CameraStream from character import MODEL_FILES, TRIANGLE_BUDGET, Character from mesh3d import MeshRenderer from pose import PoseTracker from smoothing import ScalarEMA WINDOW = "Pose Driven 3D Characters" STAGE_TOP = (58, 46, 38) STAGE_BOTTOM = (24, 20, 17) STAGE_MARGIN = 0.06 RENDER_SCALE = 1.0 # Cost is set by triangle count, not by pixels, so the quality lever cycled with # F is how finely the mesh is kept: these scale the loader's triangle budget. DETAIL_LEVELS = (1.0, 0.45, 0.2) DETAIL_LABELS = ("high", "medium", "fast") SHADOW_COLOR = (10, 8, 7) SHADOW_ALPHA = 0.55 # How fast the torso turn follows its estimate. Low, because shoulder width is # the noisiest thing the tracker reports and the turn reads as a whole-body move. TURN_SMOOTHING = 0.15 VIEWS = ("full3d", "pip_cam", "pip_3d") VIEW_LABELS = { "full3d": "3D CHARACTER", "pip_cam": "3D CHARACTER + camera inset", "pip_3d": "CAMERA + character inset", } PIP_RATIO = 0.30 MODEL_NAMES = list(MODEL_FILES.keys()) HELP_LINES_POSE = ( "POSES", " Stand so your shoulders, hips, arms and legs are visible", " The character faces you and copies your pose in real time", " Raise your right arm and its right arm goes up", " Turn your shoulders to turn its torso", " Keypoints and skeleton are drawn on the camera image", "", "KEYS", " K cycle view (3D only / cam inset / 3D inset)", " M next model N previous model", " F render quality (lower it if the frame rate drops)", " L keypoint names", " Y / P yaw / pitch character view R reset view", " D debug H help Q quit", ) class _Scaled(NamedTuple): """Keypoints resized for a smaller canvas, for drawing only.""" keypoints: np.ndarray scores: np.ndarray class PoseApp: """Body-pose driven 3D character viewer.""" def __init__(self, args: argparse.Namespace) -> None: self.args = args self.cam = CameraStream(args.camera, args.width, args.height) self.h, self.w = self.cam.shape print(f"[i] camera: {self.w}x{self.h} @ {self.cam.source_fps:.0f} FPS", flush=True) self.tracker = PoseTracker(args.pose_model, imgsz=args.pose_size, conf=args.pose_conf) print(f"[i] pose device: {self.tracker.device}", flush=True) self.renderer = MeshRenderer(self.w, self.h) self.model_index = MODEL_NAMES.index(getattr(args, "model", MODEL_NAMES[0])) self.character: Character | None = None self.yaw, self.pitch = 0.0, 0.0 self.detail_index = 0 self._load_character() self.view_index = 0 self.turn = 0.0 self.turn_s = ScalarEMA(TURN_SMOOTHING, 0.0) self._stage_cache: np.ndarray | None = None self.show_help = False self.show_debug = args.debug self.show_labels = False self.toast_text = "" self.toast_until = 0.0 self.fps = 0.0 self._last_t = time.perf_counter() self._no_pose_since: float | None = None def _load_character(self) -> None: """(Re)load the active character model.""" name = MODEL_NAMES[self.model_index] budget = int(TRIANGLE_BUDGET * DETAIL_LEVELS[self.detail_index]) self.character = Character.load(name, mirror=not self.args.no_mirror, budget=budget) self.yaw, self.pitch = self.character.default_yaw, 0.0 self.toast(f"Model: {name}") def toast(self, text: str, seconds: float = 2.0) -> None: """Show transient message.""" self.toast_text = text self.toast_until = time.time() + seconds def next_model(self, step: int) -> None: """Switch to another character model.""" self.model_index = (self.model_index + step) % len(MODEL_NAMES) self._load_character() def reset_view(self) -> None: """Back to this model's default camera angle.""" self.yaw = self.character.default_yaw if self.character else 0.0 self.pitch = 0.0 def render_character(self, canvas_shape: tuple[int, int], pose) -> np.ndarray: """Render the character posed by one detected person, or the rest pose. The software rasterizer costs time per covered pixel, so the figure is drawn at a reduced resolution and scaled up. Smooth-shaded geometry survives that far better than the frame rate survives full resolution. """ out_h, out_w = canvas_shape assert self.character is not None ch = self.character h = max(int(out_h * RENDER_SCALE), 64) w = max(int(out_w * RENDER_SCALE), 64) # Ease the torso turn toward its estimate, and back to square when the # person is lost, so a noisy shoulder measurement cannot make the figure # twitch from side to side. self.turn = self.turn_s(ch.torso_yaw(pose)) rots = ch.local_rotations(pose, torso_yaw=self.turn) if pose is not None else {} parts = ch.skin(ch.pose_globals(rots)) rot = mesh3d._rotation(self.yaw, self.pitch) distance = 3.0 mid_y, height, half_w = ch.framing(parts, rot, distance) # Fit the projected figure into the frame with margin, on both axes. margin = 1.0 - 2 * STAGE_MARGIN focal = min(h * margin / height, w * 0.5 * margin / half_w) cx = w * 0.5 cy = h * 0.5 + focal * mid_y out = self._stage(h, w) self._ground_shadow(out, parts, rot, origin=(cx, cy), focal=focal, distance=distance) out = self.renderer.render( out, parts, origin=(cx, cy), focal=focal, distance=distance, yaw=self.yaw, pitch=self.pitch, ) if (h, w) != (out_h, out_w): out = cv2.resize(out, (out_w, out_h), interpolation=cv2.INTER_LINEAR) return out def _stage(self, h: int, w: int) -> np.ndarray: """Vertical studio-gradient backdrop, built once per size.""" if self._stage_cache is None or self._stage_cache.shape[:2] != (h, w): top = np.array(STAGE_TOP, dtype=np.float32) bottom = np.array(STAGE_BOTTOM, dtype=np.float32) t = np.linspace(0.0, 1.0, h, dtype=np.float32)[:, None] column = top * (1.0 - t) + bottom * t self._stage_cache = np.repeat(column[:, None, :], w, axis=1).astype(np.uint8) return self._stage_cache.copy() def _ground_shadow(self, out: np.ndarray, parts, rot: np.ndarray, *, origin, focal, distance) -> None: """Soft elliptical contact shadow under the figure.""" cam = np.concatenate([p.positions for p in parts], axis=0) @ rot.T z = np.maximum(cam[:, 2] + distance, 1e-3) px, py = cam[:, 0] / z, cam[:, 1] / z cx, cy = origin # Sit the ellipse at the lowest projected point, spanning the silhouette. sx = int(cx + focal * float(np.median(px))) sy = int(cy - focal * float(np.percentile(py, 0.5))) rx = max(int(focal * float(np.percentile(np.abs(px), 98)) * 0.95), 6) ry = max(int(rx * 0.22), 3) # Work inside the ellipse's own bounding box: blurring the whole frame # to soften a shape this small was costing more than drawing the figure. pad = max(rx // 4, 4) x0, y0 = max(sx - rx - pad, 0), max(sy - ry - pad, 0) x1, y1 = min(sx + rx + pad, out.shape[1]), min(sy + ry + pad, out.shape[0]) if x1 - x0 < 3 or y1 - y0 < 3: return roi = out[y0:y1, x0:x1] layer = roi.copy() cv2.ellipse(layer, (sx - x0, sy - y0), (rx, ry), 0, 0, 360, SHADOW_COLOR, -1, cv2.LINE_AA) blur = max(3, (pad // 2) * 2 + 1) cv2.GaussianBlur(layer, (blur, blur), 0, dst=layer) cv2.addWeighted(layer, SHADOW_ALPHA, roi, 1.0 - SHADOW_ALPHA, 0.0, roi) def compose(self, frame: np.ndarray, poses: list) -> np.ndarray: """Build the output frame for the current view mode.""" view = VIEWS[self.view_index] pose = poses[0] if poses else None char_img = self.render_character((self.h, self.w), pose) if view == "full3d": return char_img threshold = self.character.pose_scores_threshold if self.character else 0.35 if view == "pip_cam": # The camera goes in the inset: shrink first, then draw the skeleton # at that size so its strokes stay proportional instead of vanishing. pw, ph = int(self.w * PIP_RATIO), int(self.h * PIP_RATIO) cam = cv2.resize(frame, (pw, ph)) k = np.array([pw / self.w, ph / self.h]) for p in poses: ui.draw_pose(cam, _Scaled(p.keypoints * k, p.scores), threshold=threshold) out = char_img self._blit_pip(out, cam, resize=False) else: out = frame.copy() for p in poses: ui.draw_pose(out, p, threshold=threshold, labels=self.show_labels) self._blit_pip(out, char_img) return out def _blit_pip(self, out: np.ndarray, source: np.ndarray, resize: bool = True) -> None: """Draw a picture-in-picture inset bottom-left.""" h, w = out.shape[:2] pw, ph = int(w * PIP_RATIO), int(h * PIP_RATIO) inset = cv2.resize(source, (pw, ph)) if resize else source ph, pw = inset.shape[:2] x0, y0 = 16, h - ph - 16 cv2.rectangle(out, (x0 - 3, y0 - 3), (x0 + pw + 3, y0 + ph + 3), (210, 210, 210), 2, cv2.LINE_AA) out[y0:y0 + ph, x0:x0 + pw] = inset def render(self, frame: np.ndarray, poses: list) -> np.ndarray: """Compose the full HUD frame.""" pose = poses[0] if poses else None out = self.compose(frame, poses) h, w = out.shape[:2] ui.panel(out, 0, 0, w, 60, 0.5) ui.put_text(out, VIEW_LABELS[VIEWS[self.view_index]], (16, 10), 20, (255, 255, 255)) name = MODEL_NAMES[self.model_index] status = f"{name} {'tracking' if pose is not None else 'no person detected'}" ui.put_text(out, status, (16, 38), 16, (140, 230, 140) if pose is not None else (200, 160, 100)) right = f"{self.fps:.0f} FPS {self.tracker.device.upper()}" rw = ui._render_text(right, 16, (200, 200, 200))[0].shape[1] ui.put_text(out, right, (w - rw - 16, 14), 16, (200, 200, 200)) ui.put_text(out, "H help", (w - rw - 16, 38), 15, (140, 140, 140)) if self.toast_text and time.time() < self.toast_until: tw = ui._render_text(self.toast_text, 18, (255, 255, 255))[0].shape[1] ui.panel(out, w // 2 - tw // 2 - 14, 74, tw + 28, 36, 0.65) ui.put_text(out, self.toast_text, (w // 2 - tw // 2, 82), 18, (255, 255, 255)) if self.show_debug and pose is not None: tris = sum(len(p.triangles) for p in self.character.model.primitives) ui.panel(out, 12, h - 128, 300, 108, 0.6) ui.put_text(out, f"score {pose.confidence:.2f} tris {tris}", (20, h - 120), 15, (220, 220, 220)) ui.put_text(out, f"view yaw {np.degrees(self.yaw):+.0f} pitch " f"{np.degrees(self.pitch):+.0f}", (20, h - 98), 15, (220, 220, 220)) ui.put_text(out, f"torso turn {np.degrees(self.turn):+.0f} deg", (20, h - 76), 15, (150, 230, 255)) ui.put_text(out, f"detail {DETAIL_LABELS[self.detail_index]}", (20, h - 54), 15, (220, 220, 220)) if self.show_help: ui.draw_help(out, HELP_LINES_POSE) return out def handle_key(self, key: int) -> bool: """Handle one keypress.""" if key in (ord("q"), 27): return False if key == ord("k"): self.view_index = (self.view_index + 1) % len(VIEWS) self.toast(VIEW_LABELS[VIEWS[self.view_index]]) elif key == ord("m"): self.next_model(1) elif key == ord("n"): self.next_model(-1) elif key == ord("y"): self.yaw += 0.25 elif key == ord("p"): self.pitch = float(np.clip(self.pitch + 0.2, -1.2, 1.2)) elif key == ord("r"): self.reset_view() self.toast("View reset") elif key == ord("f"): self.detail_index = (self.detail_index + 1) % len(DETAIL_LEVELS) self._load_character() self.toast(f"Detail: {DETAIL_LABELS[self.detail_index]}") elif key == ord("l"): self.show_labels = not self.show_labels self.toast("Keypoint names " + ("on" if self.show_labels else "off")) elif key == ord("d"): self.show_debug = not self.show_debug self.toast("Debug " + ("on" if self.show_debug else "off")) elif key == ord("h"): self.show_help = not self.show_help return True def run(self) -> None: """Run the main loop.""" cv2.namedWindow(WINDOW, cv2.WINDOW_NORMAL) cv2.resizeWindow(WINDOW, self.w, self.h) print("[i] window open. H help, Q quit.", flush=True) self.toast("Stand back so your full body is visible") reason = "loop finished" frames = 0 seq = -1 while True: frame, seq = self.cam.wait_next(seq) if self.cam.failed: reason = "camera stopped delivering frames" break if not self.args.no_mirror: frame = cv2.flip(frame, 1) if frame.shape[:2] != (self.h, self.w): frame = cv2.resize(frame, (self.w, self.h)) now = time.perf_counter() dt = max(now - self._last_t, 1e-4) self._last_t = now self.fps += 0.12 * (1.0 / dt - self.fps) poses = self.tracker(frame, dt) out = self.render(frame, poses) cv2.imshow(WINDOW, out) key = cv2.waitKey(1) & 0xFF frames += 1 if frames % 150 == 0: print(f"[i] frames: {frames} {self.fps:.1f} FPS people: {len(poses)}", flush=True) if key != 255 and not self.handle_key(key): reason = f"quit key pressed ({key})" break if cv2.getWindowProperty(WINDOW, cv2.WND_PROP_VISIBLE) < 1: reason = "window closed" break print(f"[i] stopped: {reason} (frames processed: {frames})", flush=True) self.cam.release() cv2.destroyAllWindows() def parse_args() -> argparse.Namespace: """Parse command line arguments.""" root = Path(__file__).resolve().parent.parent p = argparse.ArgumentParser( description="Drive a rigged 3D character with your body pose") p.add_argument("--camera", type=int, default=0) p.add_argument("--pose-model", default=str(root / "models" / "yolo26s-pose.pt")) p.add_argument("--pose-size", type=int, default=640) p.add_argument("--pose-conf", type=float, default=0.35) p.add_argument("--width", type=int, default=1280) p.add_argument("--height", type=int, default=720) p.add_argument("--model", choices=MODEL_NAMES, default=MODEL_NAMES[0], help="character to start with") p.add_argument("--no-mirror", action="store_true", help="do not flip the camera image") p.add_argument("--debug", action="store_true") return p.parse_args() def main() -> int: """Program entry point.""" args = parse_args() try: PoseApp(args).run() except (FileNotFoundError, RuntimeError, ImportError) as exc: print(f"[!] {exc}", file=sys.stderr) return 1 except KeyboardInterrupt: print("\n[i] stopped by user") return 0 if __name__ == "__main__": raise SystemExit(main())