from transformers import AutoImageProcessor, SiglipForImageClassification
from PIL import Image
import torch
import gradio as gr

# Load model and processor from HuggingFace
model_name = "prithivMLmods/Recycling-Net-11"
processor = AutoImageProcessor.from_pretrained(model_name)
model = SiglipForImageClassification.from_pretrained(model_name)

# Define recyclable and non-recyclable categories
recyclable_labels = [
    "cardboard", "glass", "metal", "paper", "plastic", "can", "carton"
]
non_recyclable_labels = [
    "food waste", "trash", "garbage", "organic"
]

# Get model class label mapping
id2label = model.config.id2label

def classify_frame(frame):
    if frame is None:
        return "No frame detected"
    
    img = Image.fromarray(frame)
    inputs = processor(images=img, return_tensors="pt")

    with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()

    pred_idx = max(range(len(probs)), key=lambda i: probs[i])
    pred_label = id2label[pred_idx].lower()

    if any(word in pred_label for word in recyclable_labels):
        return f"♻️ Recyclable ({probs[pred_idx]*100:.1f}%)"
    else:
        return f"🗑️ Non-Recyclable ({probs[pred_idx]*100:.1f}%)"

# Gradio Interface
gr.Interface(
    fn=classify_frame,
    inputs=gr.Image(source="webcam", streaming=True, label="Live Waste Feed"),
    outputs=gr.Text(label="Prediction"),
    live=True,
    title="Live Waste Classification",
    description="Classifies live webcam input into Recyclable or Non-Recyclable using 11-class model."
).launch()