import glob

import gradio as gr
import tensorflow as tf

from utils.predict import predict_label

##Create list of examples to be loaded
example_list = glob.glob("examples/*.mp4")
example_list = list(map(lambda el:[el], example_list))

demo = gr.Blocks()

with demo:
    gr.Markdown("# **<p align='center'>Video Vision Transformer on medmnist</p>**")
    
    with gr.Tab("Upload & Predict"):
        with gr.Box():
            with gr.Row():
                input_video = gr.Video(label="Input Video", show_label=True)
                output_label = gr.Label(label="Model Output", show_label=True)

        gr.Markdown("**Predict**")

        with gr.Box():
            with gr.Row():
                submit_button = gr.Button("Submit")

        gr.Markdown("The model is trained to classify videos belonging to the following classes: liver, kidney-right, kidney-left, femur-right, femur-left, bladder, heart, lung-right, lung-left, spleen and pancreas")

        gr.Examples(example_list, [input_video], [output_label], predict_label, cache_examples=True)

    submit_button.click(predict_label, inputs=input_video, outputs=output_label)

    gr.Markdown('\n Demo created by: <a href=\"https://huggingface.co/pablorodriper\"> Pablo Rodríguez</a> <br> Based on the Keras example by <a href=\"https://keras.io/examples/vision/vivit/\">Aritra Roy Gosthipaty and Ayush Thakur</a>')

demo.launch()