import gradio as gr
from transformers import pipeline
import tensorflow
import torch

global pipeline
global default_model_name
default_model_name = "google/vit-base-patch16-224"

def predict(image, model_name):
    if model_name == "":
        model_name = default_model_name
    pipe = pipeline(task="image-classification", model=model_name)
    predictions = pipe(image)
    return {p["label"]: p["score"] for p in predictions}
    
    
with gr.Blocks() as demo:
    
    with gr.Row():
        gr.Markdown(
        """
        # Settings
        [Here](https://huggingface.co/models?pipeline_tag=image-classification&sort=downloads) are some popular image classification models.  
        Or use default model **"google/vit-base-patch16-224"**
        """)
        gr.Markdown(
        """
        # Image Classifier Result
        """)
    with gr.Row():
        with gr.Column(scale=1):
            #input_model = gr.Textbox(label="Enter a custom model name:", value=default_model_name, scale=1)
            input_model = gr.Textbox(label="Enter a custom model name:", scale=1)
                
            gr.Markdown("Upload image")
            #images_input = gr.File(file_count="multiple", file_types=["image"], label="Input images", scale=1)
            #images_input = gr.Files(file_count="multiple", file_types=["image"], label="Input images", scale=1)
            input_image = gr.Image(label="Input Image", type="filepath")

        #output = gr.Label(label="Output", num_top_classes=3, scale=2)
        output = gr.Label(num_top_classes=10, scale=2)
    with gr.Row(equal_height=True):
        clear_button = gr.ClearButton(value="Clear", scale=0)
        submit_button = gr.Button(value="Submit", variant="primary", scale=0)
    submit_button.click(fn=predict, inputs=[input_image, input_model], outputs=output)
    clear_button.click(lambda: [None, None, None], outputs=[input_model, input_image, output])
    
demo.launch()