Spaces:
Sleeping
Sleeping
File size: 4,571 Bytes
ac5f9d6 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | 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.') |