from moviepy.editor import *
import face_recognition
from path import Path
from PIL import Image, ImageDraw
import numpy as np 

import gradio 
import os
os.makedirs('video', exist_ok=True)


def procss_video(video_str):
    source_frames = []
    # video_name = 'ZW9WPPMgpWP4bjjs.mp4'
    clip = VideoFileClip(video_str)
    for item in clip.iter_frames():
        source_frames.append(item)

    audioclip = clip.audio
    im2 = Image.open("mask_output.png")






    dealed_frames = []
    batch_size = 8
    for i in range(0, len(source_frames), batch_size):
        frames = source_frames[i:i+batch_size]
        # print('processing framese from {} to {}'.format(i, i+batch_size))
        batch_of_face_locations = face_recognition.batch_face_locations(frames, number_of_times_to_upsample=2, batch_size = batch_size)

        # Now let's list all the faces we found in all 128 frames
        for frame_number_in_batch, face_locations in enumerate(batch_of_face_locations):

            pil_image = Image.fromarray(frames[frame_number_in_batch])
            draw = ImageDraw.Draw(pil_image)
            # Loop through each face found in the unknown image
            for (top, right, bottom, left) in face_locations:

                im = im2.resize((int(abs(top-bottom)*0.8), int(abs(left-right)*0.8) ))
                pil_image.paste(im, (left, int((top+bottom)/2)),  im)
            dealed_frames.append(np.array(pil_image))

    output_clip = ImageSequenceClip(dealed_frames, fps=20)
    new_audioclip = CompositeAudioClip([audioclip])
    output_clip.audio = new_audioclip
    output_clip.write_videofile(os.path.join('video', 'processed_'+Path(video_str).name), fps=20,codec="libx264", audio_codec="aac")

    return os.path.join('video', 'processed_'+Path(video_str).name)
# os.mkdir('video')


def video_identity(video):
    # print(video)
    return procss_video(video)


demo = gradio.Interface(video_identity, 
                    gradio.Video(), 
                    "playable_video", 
                    examples=[
                        os.path.join(os.path.dirname(__file__), 
                                     "video/ZW9WPPMgpWP4bjjs.mp4")], 
                    cache_examples=False)

if __name__ == "__main__":
    demo.launch()