"""Validate pose detection on real anime/illustration images. Runs both YOLOv8m-pose and DWPose wholebody estimators on images from a directory and saves visualizations with keypoints overlaid + prints detected tags. Usage: python scripts/validate_pose.py [--limit N] """ import os import sys import argparse import json from pathlib import Path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import numpy as np from PIL import Image, ImageDraw, ImageOps # COCO-17 skeleton connections (for drawing pose lines) _COCO_SKELETON = [ (0, 1), (0, 2), (1, 3), (2, 4), # head (5, 6), (5, 7), (7, 9), # left arm (6, 8), (8, 10), # right arm (5, 11), (6, 12), (11, 12), # torso (11, 13), (13, 15), # left leg (12, 14), (14, 16), # right leg ] _WB_HAND_SKELETON = [ (0, 1), (1, 2), (2, 3), (3, 4), (0, 5), (5, 6), (6, 7), (7, 8), (0, 9), (9, 10), (10, 11), (11, 12), (0, 13), (13, 14), (14, 15), (15, 16), (0, 17), (17, 18), (18, 19), (19, 20), ] def _vis(kp: np.ndarray, i: int, thresh: float = 0.25) -> bool: """Check if keypoint i is visible.""" if kp.ndim != 2 or kp.shape[0] <= i: return False return kp[i, 2] >= thresh def draw_keypoints(img: Image.Image, keypoints: np.ndarray, skeleton: list, color=(255, 0, 0), radius=3, line_width=1) -> Image.Image: """Draw keypoints and skeleton lines on an image copy.""" canvas = img.convert("RGB") draw = ImageDraw.Draw(canvas) if keypoints.ndim == 2 and keypoints.shape[0] >= len(skeleton[0]) if skeleton else 17: # Draw skeleton lines for a, b in skeleton: if _vis(keypoints, a) and _vis(keypoints, b): draw.line([keypoints[a, 0], keypoints[a, 1], keypoints[b, 0], keypoints[b, 1]], fill=color, width=line_width) # Draw keypoint dots with index labels for i in range(keypoints.shape[0]): if _vis(keypoints, i): x, y = int(keypoints[i, 0]), int(keypoints[i, 1]) draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill=color) draw.text((x + 4, y - 4), str(i), fill=(255, 255, 255)) return canvas def draw_hands(img: Image.Image, keypoints: np.ndarray, hand_indices: list, colors: list) -> Image.Image: """Draw hand skeletons on an image.""" draw = ImageDraw.Draw(img) for hk, hand_color in zip(hand_indices, colors): hand_kpts = keypoints[hk:hk + 21] for a, b in _WB_HAND_SKELETON: if _vis(hand_kpts, a, 0.15) and _vis(hand_kpts, b, 0.15): draw.line([hand_kpts[a, 0], hand_kpts[a, 1], hand_kpts[b, 0], hand_kpts[b, 1]], fill=hand_color, width=1) for i in range(21): if _vis(hand_kpts, i, 0.15): x, y = int(hand_kpts[i, 0]), int(hand_kpts[i, 1]) draw.ellipse([x - 2, y - 2, x + 2, y + 2], fill=hand_color) # Draw face points (indices 23-90 in the 133-keypoint layout) face_kpts = keypoints[23:91] for i in range(face_kpts.shape[0]): if _vis(face_kpts, i, 0.15): x, y = int(face_kpts[i, 0]), int(face_kpts[i, 1]) draw.ellipse([x - 1, y - 1, x + 1, y + 1], fill=(255, 128, 0)) return img def resize_for_display(img: Image.Image, max_size=800) -> Image.Image: """Resize image for display.""" w, h = img.size if max(w, h) <= max_size: return img ratio = max_size / max(w, h) return img.resize((int(w * ratio), int(h * ratio)), Image.LANCZOS) def validate_on_image(img_path: str, output_dir: str) -> dict: """Run both pose estimators on a single image and save results + visualizations.""" result = { "image": os.path.basename(img_path), "yolo": {"detected": False, "people_count": 0, "pose_score": 0.0, "pose_tags": [], "keypoints_count": 0}, "wholebody": {"detected": False, "people_count": 0, "pose_score": 0.0, "pose_tags": [], "body_kpts": 0, "face_kpts": 0, "hand_kpts": 0}, "visualized": False, } try: img = Image.open(img_path).convert("RGB") img_exif = ImageOps.exif_transpose(img) except Exception as e: result["error"] = f"Could not load image: {e}" return result # --- Run YOLO pose (17 keypoints) --- try: from src.pose_tagger import get_pose_tagger est = get_pose_tagger() if est.ensure_loaded(): yolo_result = est.estimate(img_exif) result["yolo"] = { "detected": yolo_result.get("people_count", 0) > 0, "people_count": yolo_result.get("people_count", 0), "pose_score": float(yolo_result.get("pose_score", 0.0)), "pose_tags": yolo_result.get("pose_tags", []), "keypoints_count": len(yolo_result.get("keypoints", [])), } if yolo_result.get("keypoints"): kpts_list = yolo_result["keypoints"] if kpts_list and isinstance(kpts_list[0], (list, np.ndarray)): kpts = np.array(kpts_list[0]) if kpts.ndim == 2 and kpts.shape[0] >= 17: vis_img = draw_keypoints( resize_for_display(img_exif), kpts, _COCO_SKELETON) vis_img.save(os.path.join( output_dir, f"yolo_{os.path.basename(img_path)}")) else: result["yolo"]["error"] = "YOLO model failed to load" except Exception as e: result["yolo"]["error"] = str(e) # --- Run DWPose wholebody (133 keypoints) --- try: from src.wholebody_pose import get_wholebody_tagger wb_est = get_wholebody_tagger() if wb_est.ensure_loaded(): wb_result = wb_est.estimate(img_exif) result["wholebody"] = { "detected": wb_result.get("people_count", 0) > 0, "people_count": wb_result.get("people_count", 0), "pose_score": float(wb_result.get("pose_score", 0.0)), "pose_tags": wb_result.get("pose_tags", []), "body_kpts": len(wb_result.get("body_kpts", [])), "face_kpts": len(wb_result.get("face_kpts", [])), "hand_kpts": len(wb_result.get("hand_kpts", [])), } if wb_result.get("keypoints"): kpts = np.array(wb_result["keypoints"]) if kpts.ndim == 2 and kpts.shape[0] >= 133: # Draw body skeleton in green vis_img = draw_keypoints( resize_for_display(img_exif), kpts[:17], _COCO_SKELETON, color=(0, 255, 0), radius=4) # Draw hand skeletons + face points vis_img = draw_hands(vis_img, kpts, [91, 112], [(255, 0, 0), (0, 0, 255)]) vis_img.save(os.path.join( output_dir, f"wb_{os.path.basename(img_path)}")) result["visualized"] = True else: result["wholebody"]["error"] = "DWPose model failed to load" except Exception as e: result["wholebody"]["error"] = str(e) return result def main(): parser = argparse.ArgumentParser( description="Validate pose detection on real images") parser.add_argument("image_dir", help="Directory of images to validate") parser.add_argument("--limit", type=int, default=20, help="Max images to process") parser.add_argument("--output", default="scripts/validation_output", help="Output directory for visualizations") args = parser.parse_args() os.makedirs(args.output, exist_ok=True) # Find images (skip Neg_ prefixed negative samples) image_exts = {".jpg", ".jpeg", ".png", ".webp"} images = sorted([ str(p) for p in Path(args.image_dir).iterdir() if p.suffix.lower() in image_exts and not p.name.startswith("Neg_") ])[:args.limit] if not images: print("No images found!") return print(f"Found {len(images)} images to validate") print(f"Output directory: {args.output}") print() results = [] for i, img_path in enumerate(images): short_name = os.path.basename(img_path) print(f"[{i + 1}/{len(images)}] {short_name}") result = validate_on_image(img_path, args.output) results.append(result) # Print summary y = result["yolo"] w = result["wholebody"] yolo_err = y.get("error", "") wb_err = w.get("error", "") print(f" YOLO: detected={y['detected']}, " f"people={y['people_count']}, " f"score={y['pose_score']:.3f}, " f"tags={y['pose_tags']}") if yolo_err: print(f" YOLO ERR: {yolo_err}") print(f" WB: detected={w['detected']}, " f"people={w['people_count']}, " f"score={w['pose_score']:.3f}, " f"body_kpts={w['body_kpts']}, " f"face_kpts={w['face_kpts']}, " f"hand_kpts={w['hand_kpts']}") print(f" WB tags: {w['pose_tags']}") if wb_err: print(f" WB ERR: {wb_err}") print() # Save results summary summary_path = os.path.join(args.output, "validation_summary.json") with open(summary_path, "w", encoding="utf-8") as f: json.dump(results, f, indent=2, ensure_ascii=False) print(f"Results saved to {summary_path}") if __name__ == "__main__": main()