import gradio
import torch
from transformers import AutoImageProcessor, MobileNetV2ForImageClassification

image_processor = AutoImageProcessor.from_pretrained("Aruno/gemini-beauty")
model = MobileNetV2ForImageClassification.from_pretrained("Aruno/gemini-beauty")


def inference(img):
    inputs = image_processor(img, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs).logits[0]
        outputs = outputs.softmax(dim=0)
        outputs = {
            "attractive": outputs[0],
            "normal": outputs[1],
            "ugly": outputs[2],
            "very attractive": outputs[3],
            "very_ugly": outputs[4],
        }
    return outputs


iface = gradio.Interface(
    fn=inference,
    inputs="image",
    outputs="label",
    title="Your Attractivness",
    description="Check your attractivness",
    examples=["face_01.jpg", "face_02.jpg"],
)

iface.launch()