"""Depth + pose estimation helpers for the Depth/Pose tab. Two heavy models are lazy-loaded on first use, both inside @spaces.GPU functions so ZeroGPU schedules them as normal short-lived workloads: - Depth-Anything-V2-Small via transformers' depth-estimation pipeline - OpenposeDetector (lllyasviel/Annotators) via controlnet_aux Pose keypoints round-trip through a plain JSON-friendly list-of-dicts so they serialise cleanly into gr.State — that's what makes click-to-edit feasible. """ from __future__ import annotations import numpy as np from PIL import Image, ImageDraw, ImageEnhance import gradio as gr import torch import spaces # ── Standard OpenPose body-18 layout (COCO-18 ordering) ───────────────────── OPENPOSE_KEYPOINT_NAMES = [ "nose", "neck", "right_shoulder", "right_elbow", "right_wrist", "left_shoulder", "left_elbow", "left_wrist", "right_hip", "right_knee", "right_ankle", "left_hip", "left_knee", "left_ankle", "right_eye", "left_eye", "right_ear", "left_ear", ] # Bones as 0-indexed (a, b) pairs into OPENPOSE_KEYPOINT_NAMES. Mirrors the # OpenPose `limbSeq` constant so ControlNet/RefControl-Pose readers see the # exact skeleton topology they expect. SKELETON_EDGES = [ (1, 2), (1, 5), (2, 3), (3, 4), (5, 6), (6, 7), (1, 8), (8, 9), (9, 10), (1, 11), (11, 12), (12, 13), (1, 0), (0, 14), (14, 16), (0, 15), (15, 17), ] # OpenPose-canonical 18-joint colour table (used for both bones and dots). JOINT_COLORS = [ (255, 0, 0), (255, 85, 0), (255, 170, 0), (255, 255, 0), (170, 255, 0), (85, 255, 0), (0, 255, 0), (0, 255, 85), (0, 255, 170), (0, 255, 255), (0, 170, 255), (0, 85, 255), (0, 0, 255), (85, 0, 255), (170, 0, 255), (255, 0, 255), (255, 0, 170), (255, 0, 85), ] # ── Lazy model loaders ─────────────────────────────────────────────────────── _depth_pipeline = None _pose_detector = None def _load_depth_pipeline(): """Loaded inside the first @spaces.GPU call, then cached for the rest of the worker's lifetime. Avoids paying the cold-start cost for users who only ever use depth or only ever use pose.""" global _depth_pipeline if _depth_pipeline is None: from transformers import pipeline _depth_pipeline = pipeline( "depth-estimation", model="depth-anything/Depth-Anything-V2-Small-hf", device=0 if torch.cuda.is_available() else -1, ) return _depth_pipeline def _load_pose_detector(): global _pose_detector if _pose_detector is None: try: from controlnet_aux import OpenposeDetector except ImportError: raise gr.Error( "controlnet_aux is not installed. Add `controlnet_aux` to requirements.txt." ) _pose_detector = OpenposeDetector.from_pretrained("lllyasviel/Annotators") if torch.cuda.is_available(): _pose_detector = _pose_detector.to("cuda") return _pose_detector # ── Depth ──────────────────────────────────────────────────────────────────── @spaces.GPU def generate_depthmap(image: Image.Image) -> Image.Image: if image is None: raise gr.Error("Upload a source image first.") pipe = _load_depth_pipeline() depth = pipe(image.convert("RGB"))["depth"] # PIL L return depth.convert("RGB") # ── Pose detection → editable keypoint list ───────────────────────────────── @spaces.GPU def detect_pose(image: Image.Image) -> tuple[list[list[dict]], int, int]: """Run pose detection and return: - poses : list of dicts, one per detected person. Each pose is a list of 18 entries {"name", "x", "y", "visible"} — in *pixel* coords (not normalised), so the editor can pass click positions in directly without a coordinate transform. - w, h : source image dimensions. """ if image is None: raise gr.Error("Upload a source image first.") img = image.convert("RGB") w, h = img.size detector = _load_pose_detector() try: poses_raw = detector.detect_poses(np.array(img)) except Exception as e: raise gr.Error(f"Pose detection failed: {e}") out = [] for pose in poses_raw: body_kps = pose.body.keypoints if pose.body else [None] * 18 person = [] for i in range(18): kp = body_kps[i] if i < len(body_kps) else None name = OPENPOSE_KEYPOINT_NAMES[i] # controlnet_aux returns normalised (x, y) ∈ [0, 1] with a score. # Drop low-confidence detections so the editor starts clean. if kp is None or (getattr(kp, "score", 1.0) or 0) < 0.3: person.append({"name": name, "x": 0.0, "y": 0.0, "visible": False}) else: person.append({ "name": name, "x": float(kp.x) * w, "y": float(kp.y) * h, "visible": True, }) out.append(person) return out, w, h # ── Rendering ─────────────────────────────────────────────────────────────── def render_pose_skeleton(poses: list[list[dict]], width: int, height: int) -> Image.Image: """Skeleton on a black canvas — this is what ControlNet/RefControl-Pose wants.""" canvas = Image.new("RGB", (width, height), (0, 0, 0)) draw = ImageDraw.Draw(canvas) for pose in poses: for a, b in SKELETON_EDGES: ka, kb = pose[a], pose[b] if not (ka["visible"] and kb["visible"]): continue draw.line([ka["x"], ka["y"], kb["x"], kb["y"]], fill=JOINT_COLORS[a], width=4) for i, kp in enumerate(pose): if not kp["visible"]: continue r = 4 draw.ellipse([kp["x"] - r, kp["y"] - r, kp["x"] + r, kp["y"] + r], fill=JOINT_COLORS[i]) return canvas def render_pose_overlay( source: Image.Image, poses: list[list[dict]], active_person: int | None = None, active_joint: int | None = None, ) -> Image.Image | None: """Skeleton on top of a dimmed source — the editing view. The active joint gets a white-ringed highlight so the user can see what their next click will move.""" if source is None: return None w, h = source.size base = ImageEnhance.Brightness(source.convert("RGB")).enhance(0.45) skel = render_pose_skeleton(poses, w, h) np_base, np_skel = np.array(base), np.array(skel) mask = np_skel.any(axis=2) np_base[mask] = np_skel[mask] out = Image.fromarray(np_base) if (active_person is not None and 0 <= active_person < len(poses) and active_joint is not None and 0 <= active_joint < 18): kp = poses[active_person][active_joint] if kp["visible"]: d = ImageDraw.Draw(out) x, y = kp["x"], kp["y"] # Double ring so it stays visible on any background colour for r, col in [(10, (255, 255, 255)), (7, (0, 0, 0))]: d.ellipse([x - r, y - r, x + r, y + r], outline=col, width=2) return out # ── Edit operations ───────────────────────────────────────────────────────── def move_joint(poses, person_idx, joint_idx, x, y, width, height): """Move (and make visible) a joint. Clamped to image bounds so clicks just outside the canvas don't fly off.""" if (not poses or person_idx is None or joint_idx is None or not (0 <= person_idx < len(poses) and 0 <= joint_idx < 18)): return poses x = float(max(0, min(x, width - 1))) y = float(max(0, min(y, height - 1))) new = [list(p) for p in poses] new[person_idx][joint_idx] = { **new[person_idx][joint_idx], "x": x, "y": y, "visible": True, } return new def hide_joint(poses, person_idx, joint_idx): if (not poses or person_idx is None or joint_idx is None or not (0 <= person_idx < len(poses) and 0 <= joint_idx < 18)): return poses new = [list(p) for p in poses] new[person_idx][joint_idx] = {**new[person_idx][joint_idx], "visible": False} return new def clear_all_joints(poses): return [[{**kp, "visible": False} for kp in pose] for pose in (poses or [])] def default_pose_template(width: int, height: int) -> list[dict]: """Standing-figure skeleton centred in the canvas — handy when detection misses everyone, or when you want to draw a pose from scratch.""" cx = width / 2 s = min(width, height) * 0.40 top = height / 2 - s def y(f): return top + f * (2 * s) return [ {"name": "nose", "x": cx, "y": y(0.05), "visible": True}, {"name": "neck", "x": cx, "y": y(0.15), "visible": True}, {"name": "right_shoulder", "x": cx + s * 0.18, "y": y(0.18), "visible": True}, {"name": "right_elbow", "x": cx + s * 0.25, "y": y(0.35), "visible": True}, {"name": "right_wrist", "x": cx + s * 0.30, "y": y(0.50), "visible": True}, {"name": "left_shoulder", "x": cx - s * 0.18, "y": y(0.18), "visible": True}, {"name": "left_elbow", "x": cx - s * 0.25, "y": y(0.35), "visible": True}, {"name": "left_wrist", "x": cx - s * 0.30, "y": y(0.50), "visible": True}, {"name": "right_hip", "x": cx + s * 0.12, "y": y(0.55), "visible": True}, {"name": "right_knee", "x": cx + s * 0.13, "y": y(0.75), "visible": True}, {"name": "right_ankle", "x": cx + s * 0.14, "y": y(0.95), "visible": True}, {"name": "left_hip", "x": cx - s * 0.12, "y": y(0.55), "visible": True}, {"name": "left_knee", "x": cx - s * 0.13, "y": y(0.75), "visible": True}, {"name": "left_ankle", "x": cx - s * 0.14, "y": y(0.95), "visible": True}, {"name": "right_eye", "x": cx + s * 0.025, "y": y(0.03), "visible": True}, {"name": "left_eye", "x": cx - s * 0.025, "y": y(0.03), "visible": True}, {"name": "right_ear", "x": cx + s * 0.05, "y": y(0.05), "visible": True}, {"name": "left_ear", "x": cx - s * 0.05, "y": y(0.05), "visible": True}, ] # ── Dropdown helpers ───────────────────────────────────────────────────────── def person_choices(poses): return [f"Person {i+1}" for i in range(len(poses or []))] def parse_person_idx(label: str | None) -> int | None: if not label: return None try: return int(label.replace("Person ", "")) - 1 except ValueError: return None def joint_name_to_index(name: str | None) -> int: if not name: return -1 try: return OPENPOSE_KEYPOINT_NAMES.index(name) except ValueError: return -1