from transformers import pipeline
from PIL import Image
import gradio as gr


def get_result_by_image(input_image: gr.Image) -> dict:
    pipe = pipeline("image-classification", model="ivandrian11/vit-fruit-classifier")

    image = Image.fromarray(input_image)
    print(pipe(image))

    sorted_data = sorted(pipe(image), key=lambda x: x["score"], reverse=True)
    return sorted_data[0]


iface = gr.Interface(fn=get_result_by_image, inputs="image", outputs="json")
iface.launch()