import gradio as gr
import time
import cv2 # opencv2 package for python.
import torch
from pytube import YouTube
from PIL import Image
import numpy as np
from transformers import pipeline

segmentor = pipeline("image-segmentation", model="facebook/detr-resnet-50-panoptic")


device = 'cuda' if torch.cuda.is_available() else 'cpu'
URL = "https://www.youtube.com/watch?v=6NBwbKMyzEE" #URL to parse

def load(URL):
    
    yt = YouTube(URL)
    vid_cap = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().last().download(filename="tmp.mp4")
    global player
    player = cv2.VideoCapture(vid_cap)
    frame_num = int(player.get(cv2.CAP_PROP_POS_FRAMES))
    frame_count = int(player.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_fps = (player.get(cv2.CAP_PROP_FPS))
    tog = 0
    return vid_cap,frame_num,frame_count,frame_fps,tog
    

def fw_fn(cur,last):
    next = cur+1
    if next > last:
        next = last
    return next
def bk_fn(cur):
    next = cur-1
    if next < 0:
        next = 0
    return next
def tog_on():
    return 1,gr.Markdown.update("""<center><h7>Status: Playing 😁</h7></center>""")
def tog_off():
    return 0,gr.Markdown.update("""<center><h7>Status: Stopped 💀</h7></center>""")
        
def pl_fn(cap,cur,last,fps,pl_tog):
    player.set(cv2.CAP_PROP_POS_FRAMES, cur)
    ret, frame_bgr = player.read(cur)
    frame = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
    frame = Image.fromarray(frame)
    output = segmentor(frame)
    for i in range(len(output)):
        mask = np.array(output[i]['mask'])/255
        img = np.array(frame)
        img2 = np.zeros_like(img)
        img2[:,:,0] = mask
        img2[:,:,1] = mask
        img2[:,:,2] = mask
        seg_mix=img/img2

        #seg_out=img2.astype(np.unit8) 
        results=seg_mix.astype(np.uint8)







    
    if pl_tog ==1:
        cur+=1
    else:
        cur = cur        
    return results,cur

with gr.Blocks() as app:
    gr.Markdown("""<center><h1>Testing</h1><h4>stuff</h4></center>""")
    play_state = gr.Markdown("""<right><h7></h7></right>""")

    with gr.Row():
        with gr.Column():
            youtube_url = gr.Textbox(label="YouTube URL",value=f"{URL}")
            load_button = gr.Button("Load Video")
            output_win = gr.Video()
        with gr.Column():
            with gr.Row():
                cur_frame = gr.Number(label="Current Frame")
                fps_frames = gr.Number(label="Video FPS",interactive=False)
                total_frames = gr.Number(label="Total Frames",interactive=False)
            #run_button = gr.Button()
            with gr.Row():
                bk = gr.Button("<")
                pl = gr.Button("Play")
                st = gr.Button("Stop")
                fw = gr.Button(">")
            det_win = gr.Image(source="webcam", streaming=True)
    with gr.Row():
        pl_tog=gr.Number(visible=False)
        ins_cnt=gr.Number(visible=False)   
    pl.click(tog_on,None,[pl_tog,play_state],show_progress=False)
    st.click(tog_off,None,[pl_tog,play_state],show_progress=False)
    pl_tog.change(pl_fn,[output_win,cur_frame,total_frames,fps_frames,pl_tog],[det_win,cur_frame],show_progress=False)
    cur_frame.change(pl_fn,[output_win,cur_frame,total_frames,fps_frames,pl_tog],[det_win,cur_frame],show_progress=False)
    bk.click(bk_fn,[cur_frame],cur_frame,show_progress=False)
    fw.click(fw_fn,[cur_frame,total_frames],cur_frame,show_progress=False)
    load_button.click(load,youtube_url,[output_win,cur_frame,total_frames,fps_frames,pl_tog])
    #run_button.click(vid_play, [output_win,cur_frame], det_win) 

app.queue(concurrency_count=10).launch()