scripts/validate_pose.py

"""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 <image_dir> [--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()