import gradio as gr
import time
import os

from mhr.predict_tools import MHRStoneRecognizeMgr, MHRVedioCuter

mgr = MHRStoneRecognizeMgr(
    whole_pkl = "./whole_model.pkl",
    hole_pkl = "./hole_model.pkl",
    vedio_cutter = None,
)

def greet(v, speed_ratio, debug, progress=gr.Progress(track_tqdm=True)):
    print("v:", v)
    result = []
    ret_offset = -3 if debug else -2
    cuter = MHRVedioCuter(speed_ratio)
    for data in progress.tqdm(cuter.iter(v)):
        ok, ret = mgr.recognize_image(data)
        if not ok : continue
        print(ret[-2:])
        if len(result)>0 and result[-1][-2:] == ret[-2:]: continue
        result.append(ret[ret_offset:])
    return "".join([ str(item)+"\n" for item in result ])


with gr.Blocks() as demo:
    
    with gr.Row():
        with gr.Column():
            input_video = gr.Video()
            input_speed = gr.Slider(1, 60, value=5, label="speed_ratio")
            input_debug = gr.Checkbox(label="debug")
            btn = gr.Button(value="Submit")
            
        with gr.Column(): 
            output_res = gr.Textbox(label="result")
            
    input = [input_video, input_speed, input_debug]
    output = [output_res]

    gr.Examples(
        examples=[[os.path.join(os.path.dirname(__file__), "example/test.mp4")]],
        inputs=input, 
        outputs=output,
        fn=greet,
        cache_examples=False,
    )
    
    btn.click(greet, inputs=input, outputs=output)
    

demo.queue().launch(debug=True)