from transformers import pipeline
import gradio as gr


_MODEL_NAME = "wf"
_HF_USER = "universalml"


def prediction_function(input_file):
    # get user name of their hugging face
    model_path = _HF_USER + "/" + _MODEL_NAME
    # takes some time
    classifier = pipeline("image-classification", model=model_path)

    try:
        result = classifier(input_file)
        predictions = dict()
        labels = []
        for each_label in result:
            predictions[each_label["label"]] = each_label["score"]
            labels.append(each_label["label"])
        result = predictions
    except:
        result = "no data provided!!"

    return result


# change _MODEL_NAME parameter
def create_demo():
    demo = gr.Interface(
        fn=prediction_function,
        inputs=gr.Image(type="pil"),
        outputs=gr.Label(num_top_classes=3),
    )
    demo.launch()


create_demo()