TrashNet / app.py
valiyevfagan's picture
Update app.py
2ec297b verified
raw
history blame contribute delete
747 Bytes
from fastai.vision.all import *
import gradio as gr
# Load the exported model
learn = load_learner('trash_model(1).pkl')
# Define labels (make sure they match your model's training labels)
labels = learn.dls.vocab
# Define prediction function
def classify_trash(img):
pred_class, pred_idx, probs = learn.predict(img)
return {labels[i]: float(probs[i]) for i in range(len(labels))}
examples = ["glass.png","plastic.jpg"]
# Gradio Interface
interface = gr.Interface(
fn=classify_trash,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=5),
title="Trash Classifier",
description="Upload a trash image to classify it into one of 5 categories.",
examples=examples
)
# Launch app
interface.launch()