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

# Load the model (this happens only once when the Space starts)
model_name = "prithivMLmods/Recycling-Net-11"
processor = AutoImageProcessor.from_pretrained(model_name)
model = SiglipForImageClassification.from_pretrained(model_name)

# Mapping to Recyclable / Non-Recyclable
recyclable_classes = {
    "aluminium", "cardboard", "glass", "hard plastic", 
    "paper", "soft plastics", "takeaway cups"
}

def predict(image):
    if image is None:
        return "No image received", "0%"

    inputs = processor(images=image, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
        predicted_idx = probs.argmax(-1).item()
        confidence = probs[0][predicted_idx].item()
    
    label = model.config.id2label[predicted_idx]
    
    # Convert to binary decision
    if label.lower() in recyclable_classes:
        final = "Recyclable"
    else:
        final = "Non-Recyclable"
    
    return final, f"{confidence:.1%}", label

# Create the interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Plastic/Waste Image"),
    outputs=[
        gr.Textbox(label="Decision"),
        gr.Textbox(label="Confidence"),
        gr.Textbox(label="Original Class")
    ],
    title="Plastic Segregation - Recyclable vs Non-Recyclable",
    description="Upload one image of plastic waste"
)

demo.launch()