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")