"""Render full-body DWPose skeletons (COCO-WholeBody 133-kpt, black background) as control images.
Reuses the same DWPose inference as foot_adetailer/scripts/autolabel_dwpose.py.
NOTE: rtmlib's draw_skeleton only supports openpose_skeleton=True for 17/18/134/26 kpts; the
Wholebody 133-kpt output must be drawn with openpose_skeleton=False (train + infer on the same style)."""
import argparse, glob, os
import numpy as np


def render_one(img, pose_fn, draw_fn, out_path, kp_thr=0.3):
    import cv2    # lazy import so the module stays importable without opencv installed
    keypoints, scores = pose_fn(img)
    canvas = np.zeros_like(img)
    canvas = draw_fn(canvas, keypoints, scores, openpose_skeleton=False, kpt_thr=kp_thr)
    cv2.imwrite(str(out_path), canvas)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--imgs', required=True)
    ap.add_argument('--out', required=True)
    ap.add_argument('--kp-thr', type=float, default=0.3)
    ap.add_argument('--device', default='cuda')
    ap.add_argument('--limit', type=int, default=0)
    args = ap.parse_args()

    import cv2
    from rtmlib import Wholebody, draw_skeleton    # lazy: not in local .venv
    pose = Wholebody(to_openpose=False, mode='balanced', backend='onnxruntime', device=args.device)
    os.makedirs(args.out, exist_ok=True)
    files = sorted(sum((glob.glob(os.path.join(args.imgs, f'*.{e}')) for e in ('jpg', 'jpeg', 'png', 'webp')), []))
    if args.limit:
        files = files[:args.limit]
    for i, fp in enumerate(files):
        img = cv2.imread(fp)
        if img is None:
            continue
        stem = os.path.splitext(os.path.basename(fp))[0]
        render_one(img, lambda im: pose(im), draw_skeleton, os.path.join(args.out, stem + '.png'), args.kp_thr)
        if i % 200 == 0:
            print(f'{i}/{len(files)}')


if __name__ == '__main__':
    main()