import torch
import gradio as gr
from transformers import pipeline

# 1. Setup the AI model (downloads automatically on first run)
pipe = pipeline("image-classification", model="prithivMLmods/Augmented-Waste-Classifier-SigLIP2")

# 2. Define the classification function
def classify_image(image):
    results = pipe(image)
    
    # 1. Get the top prediction label (e.g., "plastic")
    top_label = results[0]['label']
    
    # 2. Create a "translation" dictionary for your waste types
    descriptions = {
    "Battery": "Hazardous Waste (Requires special disposal)",
    "Biological": "Biological/Food Waste (Compostable)",
    "Cardboard": "Cardboard Waste (Recyclable)",
    "Clothes": "Textile Waste (Donate or Recycle)",
    "Glass": "Glass Waste (Handle with care)",
    "Metal": "Metal Waste (Recyclable)",
    "Paper": "Paper Waste (Recyclable)",
    "Plastic": "Plastic Waste (Recyclable)",
    "Shoes": "Footwear Waste (Donate if possible)",
    "Trash": "General Trash (Non-recyclable)"
    }

    
    # 3. Get the description, or use the label if not in our list
    full_description = descriptions.get(top_label, f"Detected: {top_label}")
    
    # 4. Return the data for the UI and the App
    # We return the dictionary for the Label component and the string for the JS
    return {result["label"]: result["score"] for result in results}, full_description

js_code = """
function(result) {
    window.AppInventor.setWebViewString(result[1]);
}
"""

# 3. Create the interface
demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Label(num_top_classes=3), gr.Textbox(visible=False)], # Hidden textbox for the app
    js=js_code 
)


# 4. Launch the app
# Note: Put 'theme' here to avoid the Gradio 6.0 warning
# Set share=True to get the .gradio.live link for MIT App Inventor
demo.launch() (theme="soft")