import gradio as gr
from fastai.vision.all import load_learner, PILImage

# Re-declare custom function used during training (needed by load_learner)
def is_cat(x): 
    return x[0].isupper()

# Load your exported learner
learn = load_learner("model.pkl")

# Build a reliable mapping from raw vocab entries -> friendly class names
vocab = list(learn.dls.vocab)

# Default mapping - try to detect boolean labels first
class_name_map = {}
if all(isinstance(v, bool) for v in vocab):
    # Most likely True == Cat, False == Dog (this matches the fastbook pattern)
    class_name_map = {True: "Cat 🐱", False: "Dog 🐶"}
else:
    # If vocab contains strings (or other values), map by content
    for v in vocab:
        s = str(v).lower()
        if "cat" in s:
            class_name_map[v] = "Cat 🐱"
        elif "dog" in s:
            class_name_map[v] = "Dog 🐶"
        else:
            # fallback to string version of the vocab entry
            class_name_map[v] = str(v)

def predict(img):
    # Accept either an image file path or PIL image depending on Gradio input
    if not isinstance(img, PILImage):
        try:
            img = PILImage.create(img)
        except Exception:
            pass

    pred, pred_idx, probs = learn.predict(img)

    # raw label from the learner's vocab (might be bool)
    raw_label = learn.dls.vocab[pred_idx]

    # map to friendly string (fallback to str(raw_label))
    label = class_name_map.get(raw_label, str(raw_label))

    confidence = float(probs[pred_idx]) * 100

    return f"This is a {label} with {confidence:.2f}% confidence."

# Gradio UI (single-text output)
iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs="text")
iface.launch()