app.py

import gradio as gr
# CHANGED: We now import AutoImageProcessor instead of AutoFeatureExtractor
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# Load model from Hugging Face
MODEL_NAME = "prithivMLmods/Augmented-Waste-Classifier-SigLIP2"
device = "cuda" if torch.cuda.is_available() else "cpu"

# CHANGED: Use AutoImageProcessor
processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
model = AutoModelForImageClassification.from_pretrained(MODEL_NAME).to(device)

# Inference function
def classify_image(image):
    if image.mode != "RGB":
        image = image.convert("RGB")
    # CHANGED: Use the new 'processor' variable
    inputs = processor(images=image, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    idx = torch.argmax(probs).item()
    label = model.config.id2label[idx]
    confidence = probs[0, idx].item()
    return f"{label} ({confidence*100:.2f}%)"

# Gradio UI
title = "Waste Classifier"
description = "Upload an image of waste and classify it into categories."

interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title=title,
    description=description,
    allow_flagging="never"
)

if __name__ == "__main__":
    interface.launch()