"""Quick test script for the modified wholebody estimator."""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np
from PIL import Image

from src.wholebody_pose import get_wholebody_tagger

est = get_wholebody_tagger()
print("Loading models...")
ok = est.ensure_loaded()
print(f"Loaded: {ok}")

if ok:
    # Test with a random image
    img = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
    result = est.estimate(img)
    print(f"People detected: {result['people_count']}")
    print(f"Pose tags: {result['pose_tags']}")
    print(f"Pose score: {result['pose_score']}")
    
    # Test with a real image
    from PIL import Image
    real_img = Image.open(r"C:\Users\EvilDevolver\Downloads\NoiZE_Dataset\Без названия (1).jpg")
    result = est.estimate(real_img)
    print(f"\nReal image (Без названия (1).jpg):")
    print(f"People detected: {result['people_count']}")
    print(f"Pose tags: {result['pose_tags']}")
    print(f"Pose score: {result['pose_score']}")
    print(f"Body kpts: {len(result['body_kpts'])}")
    print(f"Face kpts: {len(result['face_kpts'])}")
    print(f"Hand kpts: {len(result['hand_kpts'])}")
else:
    print("Failed to load models!")