Update app.py

app.py

@@ -2,6 +2,9 @@ import os
 os.system("wget https://huggingface.co/akhaliq/lama/resolve/main/best.ckpt")
 os.system("pip install imageio")
 os.system("pip install albumentations==0.5.2")
+os.system("pip install opencv-python")
+os.system("pip install ffmpeg-python")
+os.system("pip install moviepy")
 import cv2
 import paddlehub as hub
 import gradio as gr
@@ -9,28 +12,100 @@ import torch
 from PIL import Image, ImageOps
 import numpy as np
 import imageio
+from moviepy.editor import *
 os.mkdir("data")
 os.rename("best.ckpt", "models/best.ckpt")
 os.mkdir("dataout")
 
+def get_frames(video_in):
+    frames = []
+    #resize the video
+    clip = VideoFileClip(video_in)
+    
+    #check fps
+    if clip.fps > 30:
+        print("vide rate is over 30, resetting to 30")
+        clip_resized = clip.resize(height=256)
+        clip_resized.write_videofile("video_resized.mp4", fps=30)
+    else:
+        print("video rate is OK")
+        clip_resized = clip.resize(height=256)
+        clip_resized.write_videofile("video_resized.mp4", fps=clip.fps)
+    
+    print("video resized to 512 height")
+    
+    # Opens the Video file with CV2
+    cap= cv2.VideoCapture("video_resized.mp4")
+    
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    print("video fps: " + str(fps))
+    i=0
+    while(cap.isOpened()):
+        ret, frame = cap.read()
+        if ret == False:
+            break
+        cv2.imwrite('kang'+str(i)+'.jpg',frame)
+        frames.append('kang'+str(i)+'.jpg')
+        i+=1
+    
+    cap.release()
+    cv2.destroyAllWindows()
+    print("broke the video into frames")
+    
+    return frames, fps
 
-def infer(img):
+def create_video(frames, fps, type):
+    print("building video result")
+    clip = ImageSequenceClip(frames, fps=fps)
+    clip.write_videofile(type + "_result.mp4", fps=fps)
     
+    return type + "_result.mp4"
+
+
+def magic_lama(img):
     
+    i = img
     img = Image.open(img)
     mask = Image.open("./masks/modelscope-mask.png")
     inverted_mask = ImageOps.invert(mask)
     
     
-    imageio.imwrite("./data/data.png", img)
-    imageio.imwrite("./data/data_mask.png", inverted_mask)
+    imageio.imwrite(f"./data/data_{i}.png", img)
+    imageio.imwrite(f"./data/data_mask_{i}.png", inverted_mask)
     os.system('python predict.py model.path=/home/user/app/ indir=/home/user/app/data/ outdir=/home/user/app/dataout/ device=cpu')
-    return "./dataout/data_mask.png", "./data/data_mask.png"
+    return f"./dataout/data_mask_{i}.png"
+
+def infer(video_in):
+    # 1. break video into frames and get FPS
+    break_vid = get_frames(video_in)
+    frames_list= break_vid[0]
+    fps = break_vid[1]
+    #n_frame = int(trim_value*fps)
+    n_frame = len(frames_list)
+    
+    if n_frame >= len(frames_list):
+        print("video is shorter than the cut value")
+        n_frame = len(frames_list)
+    
+    # 2. prepare frames result arrays
+    result_frames = []
+    print("set stop frames to: " + str(n_frame))
+    
+    for i in frames_list[0:int(n_frame)]:
+        lama_frame = magic_lama(i)
+        result_frames.append(lama_frame)
+        print("frame " + i + "/" + str(n_frame) + ": done;")
+
+    
+    final_vid = create_video(result_frames, fps, "cleaned")
+
+    files = [final_vid]
 
+    return final_vid, files
 
 inputs = [gr.Image(label="Input", source="upload", type="filepath")]
-outputs = [gr.outputs.Image(type="file", label="output"),
-           gr.outputs.Image(type="file", label="Mask")]
+outputs = [gr.Video(label="output"),
+           gr.Files(label="Download Video")]
 title = "LaMa Image Inpainting"
 description = "Gradio demo for LaMa: Resolution-robust Large Mask Inpainting with Fourier Convolutions. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below. Masks are generated by U^2net"
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2109.07161' target='_blank'>Resolution-robust Large Mask Inpainting with Fourier Convolutions</a> | <a href='https://github.com/saic-mdal/lama' target='_blank'>Github Repo</a></p>"