Spaces:
Build error
Build error
File size: 4,015 Bytes
01dcec0 |
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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
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
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 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(f"./data/data.png", img)
imageio.imwrite(f"./data/data_mask.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 f"./dataout/data_mask.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)
lama_frame = Image.open(lama_frame)
imageio.imwrite(f"cleaned_frame_{i}", lama_frame)
result_frames.append(f"cleaned_frame_{i}")
print("frame " + i + "/" + str(n_frame) + ": done;")
final_vid = create_video(result_frames, fps, "cleaned")
files = [final_vid]
return final_vid
inputs = [gr.Video(label="Input", source="upload", type="filepath")]
outputs = [gr.Video(label="output")]
title = "LaMa Video Watermark Remover"
description = "<p style='text-align: center'>LaMa: Resolution-robust Large Mask Inpainting with Fourier Convolutions. <br />This demo in meant to be used as a watermark remover on Modelscope generated videos. <br />Simply upload your modelscope video and hit Submit</p>"
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>"
examples = ["./examples/modelscope-astronaut-horse.mp4", "./examples/modelscope-panda.mp4", "./examples/modelscope-spiderman-surfing.mp4"]
gr.Interface(infer, inputs, outputs, title=title,
description=description, article=article, examples=examples).launch()
|