File size: 1,880 Bytes
e98c910 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | """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()
|