Spaces:
Sleeping
Sleeping
File size: 1,096 Bytes
44d631c 2163ed1 44d631c 2163ed1 44d631c 2163ed1 44d631c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
# app.py
import gradio as gr
import torch
from torchvision import transforms
from torchvision.models import resnet18
from PIL import Image
import numpy as np
# Class names (adjust if yours are different)
class_names = ['cardboard', 'glass', 'metal', 'paper', 'plastic', 'trash']
# Load model
model = resnet18(weights=None)
model.fc = torch.nn.Linear(model.fc.in_features, 6)
model.load_state_dict(torch.load("garbage_classifier.pt", map_location="cpu"))
model.eval()
# Image preprocessing
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
# Prediction function
def classify_image(img):
# Convert numpy.ndarray to PIL.Image
if isinstance(img, np.ndarray):
img = Image.fromarray(img)
img = transform(img).unsqueeze(0)
with torch.no_grad():
logits = model(img)
probs = torch.nn.functional.softmax(logits[0], dim=0)
return {class_names[i]: float(probs[i]) for i in range(6)}
# Launch Gradio interface
gr.Interface(fn=classify_image, inputs="image", outputs="label", title="Trash Classifier").launch()
|