Spaces:
Sleeping
Sleeping
| 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.') |