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

# 1. THE CRITICAL FIX: 
def label_func(x): return x.parent.name
__main__.label_func = label_func

# 2. LOAD THE LEARNER:
learn = load_learner('waste_model_448_final_v3.pkl')
categories = learn.dls.vocab

# 3. PREDICTION LOGIC:
def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    
    # Return a dictionary of {Category: Probability} for the Gradio UI
    return {categories[i]: float(probs[i]) for i in range(len(categories))}

# 4. GRADIO INTERFACE DESIGN:
demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="pil"), 
    outputs=gr.Label(num_top_classes=3),
    title="♻️ ConvNeXt-50 Waste Classifier",
    description="This AI model identifies waste categories with **98.65% accuracy**. Upload a clear image to begin.",
    article="Developed as part of a Mini Project focusing on high-resolution (448px) deep learning for environmental sustainability."
)

# 5. EXECUTION:
if __name__ == "__main__":
    demo.launch()