"""Control a 3D character with body pose keypoints.""" from __future__ import annotations import argparse import time 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 = "Hand Tracking Drawing" 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 AppPose: """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 = 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()