callmeuser34's picture
Create app.py
a18e6cf verified
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()