Deblur-Media / app.py
Omnibus's picture
Update app.py
6120714 verified
import gradio as gr
import yt_dlp
import os
#import json
import numpy as np
import cv2
import uuid
import moviepy
from moviepy.video.io.ffmpeg_tools import ffmpeg_extract_subclip
import moviepy.video.io.ImageSequenceClip
from moviepy.editor import *
from PIL import Image
#os.system("pip install --upgrade yt-dlp")
os.system("yt-dlp -U")
uid=uuid.uuid4()
if not os.path.exists(f'{uid}'): os.makedirs(f'{uid}')
if not os.path.exists(f'{uid}-frames'): os.makedirs(f'{uid}-frames')
if not os.path.exists(f'{uid}-rembg'): os.makedirs(f'{uid}-rembg')
esr = gr.Interface.load("spaces/Omnibus/real_esrgan_mod")
main_url= "https://omnibus-Deblur-Media.hf.space"
load_js = """
function(text_input, url_params) {
console.log(text_input, url_params);
const params = new URLSearchParams(window.location.search);
url_params = Object.fromEntries(params);
return [text_input, url_params]
}
"""
def load_video(vid):
#os.system("yt-dlp -U")
new_video_in = str(vid)
capture = cv2.VideoCapture(new_video_in)
fps = capture.get(cv2.CAP_PROP_FPS)
msc = capture.get(cv2.CAP_PROP_POS_MSEC)
frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
msc = float(frame_count/fps)
capture.release()
vid_t = float(msc)
hours = int(vid_t/360)
minutes = int(vid_t/60)-(hours*360)
seconds = vid_t-(minutes*60)-(hours*360)
vid_len = f'{hours}:{minutes}:{seconds}'
return frame_count, fps, vid_len, vid
def im_2_vid(images,fps):
this_im=cv2.imread(images[0])
height=this_im.shape[0]
width= this_im.shape[1]
size = (width, height)
movie_clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(images, fps)
movie_clip.write_videofile(f'{uid}-rembg/bg_removed-{uid}.mp4')
return (f'{uid}-rembg/bg_removed-{uid}.mp4', f'{uid}-rembg/bg_removed-{uid}.mp4')
def predict(text, url_params):
mod_url=""
mod=gr.HTML("")
out = None
valid=gr.update(visible=False)
mod_url = url_params.get('url')
return ["" + text + "", mod_url]
def dl(inp,img):
fps="Error"
out = None
out_file=[]
if img == None and inp !="":
try:
inp_out=inp.replace("https://","")
inp_out=inp_out.replace("/","_").replace(".","_").replace("=","_").replace("?","_")
if "twitter" in inp:
os.system(f'yt-dlp "{inp}" --extractor-arg "twitter:api=syndication" --trim-filenames 160 -o "{uid}/{inp_out}.mp4" -S res,mp4 --recode mp4')
else:
os.system(f'yt-dlp "{inp}" --trim-filenames 160 -o "{uid}/{inp_out}.mp4" -S res,mp4 --recode mp4')
out = f"{uid}/{inp_out}.mp4"
capture = cv2.VideoCapture(out)
fps = capture.get(cv2.CAP_PROP_FPS)
capture.release()
except Exception as e:
print(e)
out = None
elif img !=None and inp == "":
capture = cv2.VideoCapture(img)
fps = capture.get(cv2.CAP_PROP_FPS)
capture.release()
out = f"{img}"
return out,out,out,out,fps
def trim_vid(vid,start_time,end_time):
print (vid)
start_hr=float(start_time.split(":",2)[0])*360
start_min=int(start_time.split(":",2)[1])*60
start_sec=int(start_time.split(":",2)[2])
end_hr=int(end_time.split(":",2)[0])*360
end_min=int(end_time.split(":",2)[1])*60
end_sec=float(end_time.split(":",2)[2])
start=start_hr+start_min+start_sec
end=end_hr+end_min+end_sec
clip = VideoFileClip(vid)
mod_fps=clip.fps
clip = clip.subclip(start, end)
clip.write_videofile(f"{uid}-clip.mp4", fps=mod_fps)
#ffmpeg_extract_subclip(vid, start, end, targetname=f"{uid}-clip.mp4")
out= f"{uid}-clip.mp4"
capture = cv2.VideoCapture(out)
frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
capture.release()
return out,frame_count
def video_clip(program,fps,qual,model_name,denoise_strength,face_enhance,outscale):
_=None
fps=float(fps)
new_video_in = str(f"{uid}-clip.mp4")
capture = cv2.VideoCapture(new_video_in)
frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
fbox=[]
out_box=[]
cnt=0
frame_count1= int(frame_count)
for i in range(int(frame_count1)-1):
capture.set(cv2.CAP_PROP_POS_FRAMES, i)
ret, frame_f = capture.read(i)
#frame_ff = cv2.cvtColor(frame_f, cv2.COLOR_BGR2RGB)
cv2.imwrite(f'{uid}-frames/{i+1}.png',frame_f)
fbox.append(f'{uid}-frames/{i+1}.png')
mes=f'Working on {i+1} of {frame_count1}'
yield _,_,mes
yield _,_,"Frames Complete"
if qual:
mes = "Improving Quality"
yield _,_,mes
if qual:
for i,img in enumerate(fbox):
#print (img)
out = os.path.abspath(img)
out_url = f'{main_url}/file={out}'
mes=f'Improving frame {i+1} of {frame_count}'
yield _,_,mes
out = esr(out_url, model_name, float(denoise_strength), face_enhance, int(outscale))
#yield _,_,mes
#cv2.imwrite(f'{img}',out)
out_box.append(f'{out}')
print (f'{main_url}/file={out}')
#print (f'out::{out}')
print (f'out_box::{out_box}')
clip = ImageSequenceClip(out_box, fps = fps)
else:
mes="Writing Video"
clip = ImageSequenceClip(fbox, fps = fps)
yield _,_,"Writing Video"
videoclip = VideoFileClip(new_video_in)
audioclip=videoclip.audio
clip.audio=audioclip
clip.write_videofile(f"{uid}/enhanced_vid.mp4")
yield _,_,"Saving Video"
out = f"{uid}/enhanced_vid.mp4"
yield out,out,"Video Complete"
def gif_clip(program,fps,qual,model_name,denoise_strength,face_enhance,outscale):
_=None
fps=float(fps)
if program == 'imageio':
print (f"NOT adjusting fps from: {fps} to: {fps*2}")
fps=fps*2
new_video_in = str(f"{uid}-clip.mp4")
capture = cv2.VideoCapture(new_video_in)
frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
fbox=[]
out_box=[]
cnt=0
frame_count1= int(frame_count)
for i in range(int(frame_count1)-1):
capture.set(cv2.CAP_PROP_POS_FRAMES, i)
ret, frame_f = capture.read(i)
cv2.imwrite(f'{uid}-frames/{i+1}.png',frame_f)
fbox.append(f'{uid}-frames/{i+1}.png')
mes=f'Working on {i+1} of {frame_count1}'
yield _,_,mes
yield _,_,"Frames Complete"
if qual:
mes = "Improving Quality"
yield _,_,mes
if qual:
for i,img in enumerate(fbox):
out = os.path.abspath(img)
out_url = f'{main_url}/file={out}'
mes=f'Improving frame {i+1} of {frame_count}'
yield _,_,mes
out = esr(out_url, model_name, float(denoise_strength), face_enhance, int(outscale))
out_box.append(f'{out}')
print (f'{main_url}/file={out}')
clip = ImageSequenceClip(out_box, fps = fps)
else:
clip = ImageSequenceClip(fbox, fps = fps)
yield _,_,"Writing GIF"
clip.to_gif(f"{uid}/clip_gif.gif",fps=fps, program=program)
#clip.speedx(2).to_gif(f"{uid}/clip_gif.gif",program=program)
out = f"{uid}/clip_gif.gif"
yield out,out,"GIF Complete"
def improve_img(inp,model_name,denoise_strength,face_enhance,outscale):
_=None
yield _,"Reading image"
#img=cv2.imread(inp)
cv2.imwrite(f'{uid}-frames/tmp_im.png',cv2.cvtColor(inp, cv2.COLOR_BGR2RGB))
out = os.path.abspath(f'{uid}-frames/tmp_im.png')
out_url = f'{main_url}/file={out}'
yield _,"Improving image"
out = esr(out_url, model_name, float(denoise_strength), face_enhance, int(outscale))
yield out,"Complete"
def update_speed(inp,clip_speed,fps):
if "-split.mp4" in inp:
out_inp = f'{inp.split("-split.mp4",1)[0]}-split.mp4'
else:
out_inp = f'{inp.split(".mp4",1)[0]}-split.mp4'
mod_fps=float(fps)*float(clip_speed)
clip = VideoFileClip(inp)
final = moviepy.video.fx.all.speedx(clip, factor=clip_speed)
final.write_videofile(f'{out_inp}', fps=mod_fps)
out = f'{out_inp}'
return out,out,out
def echo_fn(inp):
return inp
def check_load(inp_url,outp_vid,hid_box,start_f,end_f):
if outp_vid == None and inp_url !="":
out_trim,in_vid,trim_count=trim_vid(hid_box,start_f,end_f)
elif outp_vid !=None and inp_url == "":
out_trim = None
in_vid=outp_vid
trim_count = ""
return out_trim,in_vid,trim_count
def run_samp(program,fps,qual,model_name,denoise_strength,face_enhance,outscale,frame_targ):
_=None
#fps=float(fps)
new_video_in = str(f"{uid}-clip.mp4")
capture = cv2.VideoCapture(new_video_in)
#frame_count1= int(frame_count)
i=int(frame_targ)
capture.set(cv2.CAP_PROP_POS_FRAMES, i)
ret, frame_f = capture.read(i)
#frame_ff = cv2.cvtColor(frame_f, cv2.COLOR_BGR2RGB)
cv2.imwrite(f'{uid}-frames/{i}.png',frame_f)
out = os.path.abspath(f'{uid}-frames/{i}.png')
out_url = f'{main_url}/file={out}'
yield _,"Creating Sample Frame","Creating Sample Frame"
out = esr(out_url, model_name, float(denoise_strength), face_enhance, int(outscale))
yield out,"Sample Frame Complete","Sample Frame Complete"
def new_media(inp):
if inp == "Image":
im_box = gr.update(visible=True)
vid_box = gr.update(visible=False)
gif_row = gr.update(visible=False)
vid_row = gr.update(visible=False)
#qual = gr.update(value=True,interactive=False)
return (im_box,vid_box,gif_row,vid_row)
if inp == "Video":
im_box = gr.update(visible=False)
vid_box = gr.update(visible=True)
gif_row = gr.update(visible=False)
vid_row = gr.update(visible=True)
#qual = gr.update(value=False,interactive=True)
return (im_box,vid_box,gif_row,vid_row)
if inp == "GIF":
im_box = gr.update(visible=False)
vid_box = gr.update(visible=True)
gif_row = gr.update(visible=True)
vid_row = gr.update(visible=False)
#qual = gr.update(value=False,interactive=True)
return (im_box,vid_box,gif_row,vid_row)
css='''
#component-0{
background:repeating-radial-gradient(white, cornflowerblue);
}
.padded.svelte-90oupt{
background:cornflowerblue;
}
.dark .gr-box{
background:#344d74;!important;
opacity:1;!important;
}
.p-2 {
background:#344d74;!important;
opacity:1;!important;
}
.gap-4 {
background:#6681ab;!important;
opacity:1;!important;
}
.dark .gr-padded{
background:#21314a;!important;
opacity:1;!important;
}
'''
with gr.Blocks(css=css) as app:
with gr.Row():
gr.Column()
with gr.Column():
with gr.Group():
choose_media = gr.Radio(label="Media Type",choices=["Image","Video","GIF"],value="Image")
with gr.Accordion("Quality Options", open=False):
with gr.Row():
model_name = gr.Dropdown(label="Real-ESRGAN Model",
choices=["RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B",
"RealESRGAN_x2plus", "realesr-general-x4v3"],
value="realesr-general-x4v3", show_label=True)
face_enhance = gr.Checkbox(label="Face Enhancement using GFPGAN",
value=False, show_label=True)
denoise_strength = gr.Slider(label="Denoise Strength",
minimum=0, maximum=1, step=0.1, value=0.5)
outscale = gr.Slider(label="Image Upscaling Factor",
minimum=1, maximum=10, step=1, value=1, show_label=True)
with gr.Box(visible=True) as im_box:
with gr.Group():
inp_img_url=gr.Textbox(label="Image URL")
load_im_btn=gr.Button("Load URL")
inp_img=gr.Image(label = "Input Image")
with gr.Row():
proc_im_btn=gr.Button("Improve")
im_stat=gr.Textbox(label="Status")
outp_img=gr.Image(label="Enhanced Image")
with gr.Box(visible=False) as vid_box:
with gr.Group():
inp_url = gr.Textbox(label="Video URL")
go_btn = gr.Button("Run")
outp_vid=gr.Video(format="mp4")
with gr.Row():
frame_count=gr.Textbox(label="Frame Count",interactive = False)
fps=gr.Textbox(label="FPS",interactive = False)
outp_file=gr.Files()
clip_speed = gr.Slider(label="Speed", minimum=0.01, maximum=2, value=1, step=0.01)
speed_btn = gr.Button("Update Speed")
with gr.Row():
start_f = gr.Textbox(label = "Start", value = "0:00:00", placeholder = "0:00:23",interactive = True)
end_f = gr.Textbox(label = "End", value = "0:00:05", placeholder = "0:00:54",interactive = True)
trim_count = gr.Textbox(label="Trimmed Frames")
trim_btn=gr.Button("Trim")
out_trim=gr.Video(format="mp4")
hid_box = gr.Textbox(visible=False)
hid_fps = gr.Textbox(visible=False)
with gr.Row():
qual=gr.Checkbox(label="Improve Quality",value=False)
with gr.Accordion("Sample Frame",open=False):
with gr.Row():
frame_targ=gr.Number(label="Frame",value=1)
samp_btn=gr.Button("Get Sample")
samp_im = gr.Image(label="Sample Frame")
with gr.Column(visible=False) as gif_row:
with gr.Row():
choose_prog=gr.Radio(label="GIF Engine",choices=['imageio','ImageMagick','ffmpeg'], value = "imageio")
with gr.Row():
gif_btn = gr.Button("Make GIF")
gif_stat=gr.Textbox(label="Status")
with gr.Row(visible=False) as vid_row:
vid_btn = gr.Button("Make Video")
vid_stat=gr.Textbox(label="Status")
with gr.Row():
with gr.Column():
gif_show = gr.Video()
with gr.Column():
gif_file = gr.Files()
gr.Column()
with gr.Row(visible=False):
text_input=gr.Textbox()
text_output=gr.Textbox()
url_params=gr.JSON()
def load_im_fn(inp):
return inp
load_im_btn.click(load_im_fn,inp_img_url,inp_img)
proc_im_btn.click(improve_img,[inp_img,model_name,denoise_strength,face_enhance,outscale],[outp_img,im_stat])
choose_media.change(new_media,choose_media,[im_box,vid_box,gif_row,vid_row])
samp_btn.click(run_samp,[choose_prog,fps,qual,model_name,denoise_strength,face_enhance,outscale,frame_targ],[samp_im,gif_stat,vid_stat])
speed_btn.click(update_speed,[hid_box,clip_speed,hid_fps],[outp_vid,hid_box,outp_file])
gif_btn.click(gif_clip,[choose_prog,fps,qual,model_name,denoise_strength,face_enhance,outscale],[gif_show,gif_file,gif_stat])
vid_btn.click(video_clip,[choose_prog,fps,qual,model_name,denoise_strength,face_enhance,outscale],[gif_show,gif_file,vid_stat])
trim_btn.click(trim_vid,[hid_box,start_f,end_f],[out_trim,trim_count])
outp_vid.change(load_video,outp_vid,[frame_count,fps,end_f,hid_box]).then(trim_vid,[hid_box,start_f,end_f],[out_trim,trim_count])
go_btn.click(dl,[inp_url,outp_vid],[outp_vid,outp_file,out_trim,hid_box,hid_fps])
app.load(fn=predict, inputs=[text_input,url_params], outputs=[text_output,text_input],_js=load_js)
app.queue().launch(show_api=False)