app.py

import gradio as gr 
import yt_dlp
import os 
import json
from rembg import remove as rm
import cv2
import uuid
uid=uuid.uuid4()
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')




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 rem_cv(inp):
    cap = cv2.VideoCapture(f'{inp}')
    
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3))
    fgbg = cv2.bgsegm.createBackgroundSubtractorGMG()
    
    while True:
        ret, frame = cap.read()
        
        fgmask = fgbg.apply(frame)
        fgmask = cv2.morphologyEx(fgmask, cv2.MORPH_OPEN, kernel)
        yield (fgmask)

def get_rate(vid_in):
    cap = cv2.VideoCapture(vid_in)

    while(cap.isOpened()): 
      
        # Capture frame-by-frame 
      
        ret, frame = cap.read() 
      
        # if video finished or no Video Input 
        if not ret: 
            break
      
        # Our operations on the frame come here 
        gray = frame 
      
        # resizing the frame size according to our need 
        gray = cv2.resize(gray, (500, 300)) 
      
        # font which we will be using to display FPS 
        font = cv2.FONT_HERSHEY_SIMPLEX 
        # time when we finish processing for this frame 
        new_frame_time = time.time() 
      
        # Calculating the fps 
      
        # fps will be number of frame processed in given time frame 
        # since their will be most of time error of 0.001 second 
        # we will be subtracting it to get more accurate result 
        fps = 1/(new_frame_time-prev_frame_time) 
        prev_frame_time = new_frame_time 


def load_video(vid):
    new_video_in = str(vid)
    capture = cv2.VideoCapture(new_video_in)
    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    tmp_frame='1'
    capture.release()



    
    return frame_count 



def capture_function(vid=None,frame_count=None):
    new_video_in = str(vid)
    capture = cv2.VideoCapture(new_video_in)
    #frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    
    fbox=[]
    cnt=0
    frame_count1= int(frame_count)
    for i in range(int(frame_count1)):
        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}/{i}.png',frame_ff)
        fbox.append(f'{uid}/{i}.png')
        
    #capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_count)-1)
    #ret, frame_f = capture.read()
    #frame_ff = cv2.cvtColor(frame_f, cv2.COLOR_BGR2RGB)
    #cv2.imwrite(f'{uid}/{int(frame_count)-1}.png',frame_ff)
    #fbox.append(f'{uid}/{int(frame_count)-1}.png')
    #capture.release()
    return fbox,fbox

    
def rem_bg(vid,frame_count):
    new_video_in = str(vid)
    capture = cv2.VideoCapture(new_video_in)    
    fbox2=[]
    cnt=0
    frame_count1= int(frame_count)
    for i in range(int(frame_count1)):
        capture.set(cv2.CAP_PROP_POS_FRAMES, i)
        ret, frame_f = capture.read(i)
        frame_ff = cv2.cvtColor(frame_f, cv2.COLOR_BGR2RGB)
        out = rm(frame_ff)
        print (i)
        cv2.imwrite(f'{uid}-rembg/{i}.png',out)
        fbox2.append(f'{uid}-rembg/{i}.png')
        
    #capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_count)-1)
    #ret, frame_f = capture.read()
    #frame_ff = cv2.cvtColor(frame_f, cv2.COLOR_BGR2RGB)
    #cv2.imwrite(f'{uid}/{int(frame_count)-1}.png',frame_ff)
    #fbox.append(f'{uid}/{int(frame_count)-1}.png')
    #capture.release()
    #out = cv2.VideoWriter('video.avi',cv2.VideoWriter_fourcc(*'DIVX'), 15, size)
    return fbox2

    


    
def rem_bg_og(inp):
    #transparent background .mov
    #os.system(f'backgroundremover -i "{inp}" -tv -o "{inp}.mov"')
    #video video to be overlayed output
    #os.system(f' backgroundremover -i "/path/to/video.mp4" -tov "/path/to/videtobeoverlayed.mp4" -o "output.mov"')
    #video over image file 
    #os.system(f'backgroundremover -i "/path/to/video.mp4" -toi "/path/to/videtobeoverlayed.mp4" -o "output.mov"')
    #output to transparent GIF
    #os.system(f'backgroundremover -i "/path/to/video.mp4" -tg -o "output.gif"')
    #output to matte background
    os.system(f'backgroundremover -i "{inp}" -mk -o "{inp}.matte.mp4"')

    
    return f'{inp}.matte.mp4'
def predict(text, url_params):
    mod_url=""
    mod=gr.HTML("")
    out = None
    valid=gr.update(visible=False)
    mod_url = url_params.get('url')
    print (mod_url)
    return ["" + text + "", mod_url]

def dl(inp):
    out = None
    out_file=[]
    try:
        inp_out=inp.replace("https://","")
        inp_out=inp_out.replace("/","_").replace(".","_")
        os.system(f'yt-dlp "{inp}" --trim-filenames 100 -o "{inp_out}.mp4"')  
        out = f"{inp_out}.mp4"
    except Exception as e:
        print (e)
        out = None
    return out,out,out

with gr.Blocks() as app:
    with gr.Tab("Load"):
        inp_url = gr.Textbox()
        go_btn = gr.Button("Run")
        with gr.Row():
            with gr.Column():
                outp_vid=gr.Video()
            with gr.Column():
                frame_count=gr.Textbox()
                outp_file=gr.Files()
    with gr.Tab("Frames"):
        frame_btn = gr.Button("Get Frames")
        with gr.Row():
            with gr.Column():
                frame_gal = gr.Gallery(columns=6)
            with gr.Column():
                frame_file = gr.Files()
        
    with gr.Tab("Rem BG"):
        with gr.Row():    
            with gr.Column():
                rem_btn=gr.Button()
                in_vid=gr.Video()
            with gr.Column():
                #rem_vid=gr.Video()
                rem_vid=gr.Gallery(columns=6)
    with gr.Row(visible=False):
        text_input=gr.Textbox()
        text_output=gr.Textbox()
        url_params=gr.JSON()
    outp_vid.change(load_video,outp_vid,frame_count)
    frame_btn.click(capture_function,[outp_vid,frame_count],[frame_gal,frame_file])
    rem_btn.click(rem_bg,[outp_vid,frame_count],rem_vid)
    go_btn.click(dl,inp_url,[outp_vid,in_vid,outp_file])
    app.load(fn=predict, inputs=[text_input,url_params], outputs=[text_output,text_input],_js=load_js)
app.launch()