Spaces:
Paused
Paused
| from easy_dwpose import DWposeDetector | |
| from PIL import Image | |
| import cv2 | |
| import os | |
| dwpose = DWposeDetector(device="cpu") | |
| def get_pose(img, outfile): | |
| #pil_image = Image.open("imgs/"+img).convert("RGB") | |
| #skeleton = dwpose(pil_image, output_type="np", include_hands=True, include_face=False) | |
| #img.thumbnail((512,512)) | |
| out_img = dwpose(img, include_hands=True, include_face=True) | |
| #print(pose['bodies']) | |
| out_img.save("imgs/"+outfile) | |
| return out_img | |
| def extract_frames(video_path, fps): | |
| video_capture = cv2.VideoCapture(video_path) | |
| frame_count = 0 | |
| frames = [] | |
| fps_in = video_capture.get(cv2.CAP_PROP_FPS) | |
| fps_out = fps | |
| index_in = -1 | |
| index_out = -1 | |
| while True: | |
| success = video_capture.grab() | |
| if not success: break | |
| index_in += 1 | |
| out_due = int(index_in / fps_in * fps_out) | |
| if out_due > index_out: | |
| success, frame = video_capture.retrieve() | |
| if not success: | |
| break | |
| index_out += 1 | |
| if debug: | |
| frame_filename = os.path.join('imgs', "frame_"+str(frame_count)+".png") | |
| cv2.imwrite(frame_filename, frame) | |
| frame_count += 1 | |
| frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) | |
| video_capture.release() | |
| print(f"Extracted {frame_count} frames") | |
| return frames | |
| frames = extract_frames("imgs/ballet.mp4", 12) | |
| target_poses = [] | |
| for i, f in enumerate(frames): | |
| get_pose(f, "out/"+str(i)+".png") | |