import gradio as gr
from fastai.vision.all import *

# Load the trained model
learn = load_learner('clocker.pkl')

def guess_if_she_is_trans(img):
    # Predict if the woman is transgender
    pred, _, probs = learn.predict(img)
    
    # Create definitions for gender probability
    transgender_probs = {learn.dls.vocab[i]: float(probs[i]) for i in range(len(learn.dls.vocab))}
    
    # Return the predicted gender and probabilities
    return pred, transgender_probs

# Create the Gradio interface
demo = gr.Interface(
    fn=guess_if_she_is_trans,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Label(num_top_classes=1, label="My guess..."),
        gr.Label(num_top_classes=5, label="Transfem probability")
    ],
    examples=[
        ["average_woman.jpg"],
        ["transgender_woman.jpg"]
    ],
    title="Transfem Clocker",
    description="Upload a photo of a woman and this will guess if she's trans."
)

# Launch the interface
demo.launch()