lines = [] lines.append("import sys, cv2, numpy as np, json, os, glob, torch") lines.append("sys.path.insert(0, '.')") lines.append("from annotator.dwpose import DWposeDetector") lines.append("from annotator.dwpose.wholebody import Wholebody") lines.append("") lines.append("BODY_KPS = list(range(0, 18))") lines.append("FACE_KPS = list(range(24, 92))") lines.append("LHAND_KPS = list(range(92, 113))") lines.append("RHAND_KPS = list(range(113, 134))") lines.append("NAMES = ['nose','neck','r_sho','r_elb','r_wri','l_sho','l_elb','l_wri',") lines.append(" 'r_hip','r_kne','r_ank','l_hip','l_kne','l_ank','r_eye','l_eye','r_ear','l_ear']") lines.append("") lines.append("def run(video_path, out_dir):") lines.append(" os.makedirs(out_dir, exist_ok=True)") lines.append(" name = os.path.splitext(os.path.basename(video_path))[0]") lines.append(" print('Processing:', name)") lines.append(" det = DWposeDetector()") lines.append(" wb = det.pose_estimation") lines.append(" cap = cv2.VideoCapture(video_path)") lines.append(" fps = cap.get(cv2.CAP_PROP_FPS)") lines.append(" W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))") lines.append(" H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))") lines.append(" tot = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))") lines.append(" print(W, 'x', H, '|', fps, 'fps |', tot, 'frames')") lines.append(" writer = cv2.VideoWriter(os.path.join(out_dir, name+'_skeleton.mp4'),") lines.append(" cv2.VideoWriter_fourcc(*'mp4v'), fps, (W,H))") lines.append(" kps_all, sc_all, frames = [], [], []") lines.append(" i = 0") lines.append(" while True:") lines.append(" ret, frame = cap.read()") lines.append(" if not ret: break") lines.append(" with torch.no_grad():") lines.append(" drawn = det(frame)") lines.append(" kps, scores = wb(frame)") lines.append(" writer.write(drawn)") lines.append(" kp = kps[0]; sc = scores[0]") lines.append(" kps_all.append(kp.copy()); sc_all.append(sc.copy())") lines.append(" frames.append({'frame':i,") lines.append(" 'body':{'names':NAMES,") lines.append(" 'coords_px':[kp[j].tolist() for j in BODY_KPS],") lines.append(" 'coords_norm':[[round(kp[j][0]/W,4),round(kp[j][1]/H,4)] for j in BODY_KPS],") lines.append(" 'scores':[round(float(sc[j]),3) for j in BODY_KPS]},") lines.append(" 'left_hand':{'coords_px':[kp[j].tolist() for j in LHAND_KPS],") lines.append(" 'coords_norm':[[round(kp[j][0]/W,4),round(kp[j][1]/H,4)] for j in LHAND_KPS],") lines.append(" 'scores':[round(float(sc[j]),3) for j in LHAND_KPS]},") lines.append(" 'right_hand':{'coords_px':[kp[j].tolist() for j in RHAND_KPS],") lines.append(" 'coords_norm':[[round(kp[j][0]/W,4),round(kp[j][1]/H,4)] for j in RHAND_KPS],") lines.append(" 'scores':[round(float(sc[j]),3) for j in RHAND_KPS]},") lines.append(" 'face':{'coords_px':[kp[j].tolist() for j in FACE_KPS],") lines.append(" 'coords_norm':[[round(kp[j][0]/W,4),round(kp[j][1]/H,4)] for j in FACE_KPS],") lines.append(" 'scores':[round(float(sc[j]),3) for j in FACE_KPS]}})") lines.append(" i += 1") lines.append(" if i%30==0 or i==tot: print(' ['+str(i)+'/'+str(tot)+']')") lines.append(" cap.release(); writer.release()") lines.append(" with open(os.path.join(out_dir,name+'_keypoints.json'),'w') as f:") lines.append(" json.dump({'video':name,'fps':fps,'width':W,'height':H,'frames':frames},f,indent=2)") lines.append(" A = np.array(kps_all); B = np.array(sc_all)") lines.append(" np.save(os.path.join(out_dir,name+'_kps.npy'), A)") lines.append(" np.save(os.path.join(out_dir,name+'_scores.npy'), B)") lines.append(" print(' Done. kps shape:', A.shape)") lines.append("") lines.append("if __name__=='__main__':") lines.append(" IN = r'C:\\Users\\harsh\\DWPose\\input_videos'") lines.append(" OUT = r'C:\\Users\\harsh\\DWPose\\output_poses'") lines.append(" vids = sorted(glob.glob(os.path.join(IN,'*.mp4')))") lines.append(" print('Found',len(vids),'videos')") lines.append(" for v in vids: run(v, OUT)") lines.append(" print('ALL DONE')") with open('extract_pose.py', 'w', encoding='utf-8') as f: f.write('\n'.join(lines)) print('Done! extract_pose.py written.')