from fastai.vision.all import *
import gradio as gr
from PIL import Image
import numpy as np

# Load the model
learn = load_learner('export.pkl')
categories = learn.dls.vocab

# Create a mapping from original label → display label
pretty_labels = {cat: cat.replace('_', ' ').title() for cat in categories}

# Display-friendly veggie list
veggie_list = ', '.join(pretty_labels[cat] for cat in categories)
description = f"Upload a vegetable image to classify it.\n\n The model can recognize: {veggie_list}."


# Prediction function
def classify_image(img):
    if isinstance(img, str):  # If it's a filepath
        img = Image.open(img)
    elif isinstance(img, np.ndarray):  # If it's a NumPy array
        img = Image.fromarray(img)
    # If it's already a PIL Image, no conversion needed

    pred, idx, probs = learn.predict(img)
    return {
        pretty_labels[c]: float(probs[i])
        for i, c in enumerate(categories)
    }

# Gradio UI
image = gr.Image(type="pil", label="Upload a vegetable image")
label = gr.Label(label="Prediction")

intf = gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
    title="Vegetable Classifier",
    description=description
)

intf.launch()