app.py

import gradio as gr
import random
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from transformers import CLIPProcessor, CLIPModel



# ------------------ LOAD CLASSIFIER ------------------
clip_model_name = "openai/clip-vit-base-patch32"

clip_model = CLIPModel.from_pretrained(clip_model_name)
clip_processor = CLIPProcessor.from_pretrained(clip_model_name)

clip_model.eval()

CLIP_LABELS = {
    "paper/cardboard": [
        "a photo of cardboard packaging waste",
        "a photo of paper waste or newspaper"
    ],
    "plastic": [
        "a photo of plastic bottle or plastic packaging waste"
    ],
    "glass": [
        "a photo of glass bottle or glass container waste"
    ],
    "metal": [
        "a photo of metal can or metal packaging waste"
    ],
    "organic waste": [
        "a photo of food waste or leftovers",
        "a photo of organic waste such as fruit peels or leaves"
    ],
    "hazardous waste": [
        "a photo of a used battery waste item",
        "a photo of hazardous household waste"
    ],
    "electronic waste": [
        "a photo of electronic waste such as cables, chargers, or devices"
    ],
    "textile": [
        "a photo of textile waste or old clothing"
    ],
    "general trash": [
        "a photo of mixed general trash or landfill waste"
    ]
}


def classify_image(image):
    all_prompts = []
    prompt_to_category = {}

    for category, prompts in CLIP_LABELS.items():
        for p in prompts:
            all_prompts.append(p)
            prompt_to_category[p] = category

    inputs = clip_processor(
        text=all_prompts,
        images=image,
        return_tensors="pt",
        padding=True
    )

    with torch.no_grad():
        outputs = clip_model(**inputs)

    logits = outputs.logits_per_image[0]
    probs = logits.softmax(dim=0)

    # merge probabilities by category
    category_scores = {}
    for i, p in enumerate(all_prompts):
        cat = prompt_to_category[p]
        category_scores[cat] = category_scores.get(cat, 0) + float(probs[i])
    total = sum(category_scores.values())
    for k in category_scores:
        category_scores[k] /= total

    # Return top 3
    top3 = sorted(
        category_scores.items(),
        key=lambda x: x[1],
        reverse=True
    )[:3]

    return dict(top3)


# ------------------ LOAD CHAT MODEL ------------------
tiny_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

tokenizer = AutoTokenizer.from_pretrained(tiny_model)
chat_model = AutoModelForCausalLM.from_pretrained(
    tiny_model,
    dtype=torch.bfloat16, 
    device_map="auto",
    low_cpu_mem_usage=True)

pipe = pipeline(
    "text-generation",
    model=chat_model,
    tokenizer=tokenizer,
    device_map="auto",
    max_new_tokens=7
)

def clean_chat_output(full_text):

    if "<|assistant|>" in full_text:
        full_text = full_text.split("<|assistant|>")[-1]

    lines = full_text.strip().split("\n")
    if lines[0].lower().startswith("item:"):
        lines = lines[1:]

    return "\n".join(lines).strip()

def explain_recycling(class_label):
    system_msg = {
        "role": "system",
        "content": (
            "You are an expert in waste sorting. "
            "You ALWAYS answer using exactly two bullet points:\n"
            "• Recycling type: <Item category>\n"
            "• Disposal: <clear, detailed correct sentence>\n"
            "No extra text, no introductions, no explanations."
        )
    }
    user_msg = {
        "role": "user",
        "content": f"Item: {class_label}\nReturn the two bullet points now."
    }
    messages = [system_msg, user_msg]
    
    prompt = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )

    outputs = pipe(
        prompt,
        max_new_tokens=7,
        do_sample=True,
        top_p = 0.9,
        temperature=0.3
    )

    raw = outputs[0]["generated_text"]
    return clean_chat_output(raw)


# ------------------ PIPELINE ------------------
waste_analyzation_v1 = {
    "paper/cardboard": {
        "Recycling type": "Paper & cardboard recycling",
        "Disposal": "Reuse clean paper for notes, wrapping, or simple DIY crafts; flatten remaining cardboard and recycle it.",
        "Tips": "Wet or food-stained paper should be composted or thrown away.",
        "Extra": "Loose paper should be stacked to prevent it from blowing away during collection."
    },

    "plastic": {
        "Recycling type": "Plastic recycling",
        "Disposal": "Reuse sturdy plastic containers for storage or organization; rinse and recycle those you no longer need.",
        "Tips": "Do not recycle soft plastics unless your area has a dedicated drop-off.",
        "Extra": "Residue left inside containers can contaminate other recyclables."
    },

    "glass": {
        "Recycling type": "Glass recycling",
        "Disposal": "Reuse jars for food storage, plants, or small items; recycle cracked or unused glass containers.",
        "Tips": "Wrap broken glass separately if your program does not accept it.",
        "Extra": "Glass recycling reduces the need for sand mining."
    },

    "metal": {
        "Recycling type": "Metal recycling",
        "Disposal": "Reuse metal tins for storage or tools; rinse and recycle cans that are no longer useful.",
        "Tips": "Crushing cans is usually allowed, but check local rules first.",
        "Extra": "Metal can be recycled repeatedly without losing strength."
    },

    "organic waste": {
        "Recycling type": "Organic waste composting",
        "Disposal": "Turn food scraps into compost for plants or soil; otherwise place them in an organic waste bin.",
        "Tips": "Drain excess liquids to reduce odor and pests.",
        "Extra": "Compost improves soil structure and water retention."
    },

    "hazardous waste": {
        "Recycling type": "Hazardous waste disposal",
        "Disposal": "Bring batteries and chemicals to designated hazardous waste collection sites.",
        "Tips": "Never place batteries loose in household trash.",
        "Extra": "Battery fires are a leading cause of recycling facility damage."
    },

    "electronic waste": {
        "Recycling type": "Electronic waste recycling",
        "Disposal": "Take electronics to certified e-waste recycling centers.",
        "Tips": "Remove SIM cards, memory cards, and batteries before recycling.",
        "Extra": "E-waste contains materials that are harmful if landfilled."
    },

    "textile": {
        "Recycling type": "Textile recycling",
        "Disposal": "Reuse old clothes as cleaning rags or donate wearable items; recycle damaged textiles if possible.",
        "Tips": "Bag textiles to keep them dry and reusable.",
        "Extra": "Textiles should never be placed in regular recycling bins."
    },

    "general trash": {
        "Recycling type": "General waste",
        "Disposal": "Throw non-recyclable or contaminated items into the general trash bin.",
        "Tips": "Check for mixed materials before assuming an item is recyclable.",
        "Extra": "Trash contamination is a major reason recyclables are rejected."
    }
}

waste_analyzation_v2 = {
    "paper/cardboard": {
        "Recycling type": "Paper recycling",
        "Disposal": "Reuse clean paper for crafts, packaging, or notes; recycle what you cannot reuse.",
        "Tips": "Remove food residue, plastic coatings, or heavy grease.",
        "Extra": "Paper fibers weaken each time they are recycled."
    },

    "plastic": {
        "Recycling type": "Plastic recycling",
        "Disposal": "Repurpose containers for storage or household use, then recycle accepted plastics.",
        "Tips": "Check the recycling number if unsure whether a plastic is accepted.",
        "Extra": "Most recycling programs only accept rigid plastics."
    },

    "glass": {
        "Recycling type": "Glass recycling",
        "Disposal": "Reuse jars for food, tools, or décor; recycle bottles and jars when no longer needed.",
        "Tips": "Do not mix ceramics or heat-resistant glass with bottles.",
        "Extra": "Glass recycling saves energy and raw materials."
    },

    "metal": {
        "Recycling type": "Metal recycling",
        "Disposal": "Reuse metal containers where possible; recycle clean aluminum and steel cans.",
        "Tips": "Sharp metal edges should be handled carefully.",
        "Extra": "Metal recycling requires far less energy than mining."
    },

    "organic waste": {
        "Recycling type": "Compostable waste",
        "Disposal": "Compost food and garden waste to produce natural fertilizer.",
        "Tips": "Avoid adding bones or cooked food unless allowed.",
        "Extra": "Composting reduces methane from landfills."
    },

    "hazardous waste": {
        "Recycling type": "Hazardous material handling",
        "Disposal": "Dispose of chemicals, batteries, and paints at approved collection facilities.",
        "Tips": "Keep hazardous waste in original containers if possible.",
        "Extra": "Improper disposal can contaminate soil and water."
    },

    "electronic waste": {
        "Recycling type": "E-waste recycling",
        "Disposal": "Recycle electronics through official e-waste programs or retailers.",
        "Tips": "Back up and erase data before disposal.",
        "Extra": "Many electronics contain toxic metals."
    },

    "textile": {
        "Recycling type": "Clothing reuse & recycling",
        "Disposal": "Turn old textiles into rags or donate wearable clothes before recycling.",
        "Tips": "Repairing clothes can delay disposal.",
        "Extra": "Textile waste is growing rapidly worldwide."
    },

    "general trash": {
        "Recycling type": "Landfill waste",
        "Disposal": "Dispose of items that cannot be recycled or composted.",
        "Tips": "When in doubt, trash is better than contaminating recycling.",
        "Extra": "One wrong item can ruin an entire recycling load."
    }
}

waste_analyzation_v3 = {
    "paper/cardboard": {
        "Recycling type": "Paper & cardboard recycling",
        "Disposal": "Reuse clean paper for simple DIY or notes, then recycle the rest.",
        "Tips": "Trash paper that is soaked or heavily stained.",
        "Extra": "Clean paper improves recycling efficiency."
    },

    "plastic": {
        "Recycling type": "Plastic recycling",
        "Disposal": "Reuse containers for storage and recycle accepted plastics afterward.",
        "Tips": "Caps should usually stay on bottles.",
        "Extra": "Reducing plastic use is often more effective than recycling."
    },

    "glass": {
        "Recycling type": "Glass recycling",
        "Disposal": "Reuse glass jars when possible, recycle broken or unused containers.",
        "Tips": "Broken glass may be rejected by sorting machines.",
        "Extra": "Glass can be recycled endlessly."
    },

    "metal": {
        "Recycling type": "Metal recycling",
        "Disposal": "Reuse metal tins or recycle clean metal items.",
        "Tips": "Ensure containers are empty before recycling.",
        "Extra": "Recycled metal saves energy and resources."
    },

    "organic waste": {
        "Recycling type": "Organic waste composting",
        "Disposal": "Convert organic waste into compost for plants or gardens.",
        "Tips": "Keep organic waste sealed to avoid pests.",
        "Extra": "Finished compost can enrich garden soil."
    },

    "hazardous waste": {
        "Recycling type": "Hazardous waste safety",
        "Disposal": "Take hazardous items to proper disposal facilities.",
        "Tips": "Never mix hazardous waste with recyclables.",
        "Extra": "Safe disposal prevents fires and pollution."
    },

    "electronic waste": {
        "Recycling type": "Electronic recycling",
        "Disposal": "Recycle electronics instead of throwing them away.",
        "Tips": "Check for local e-waste events.",
        "Extra": "E-waste recycling reduces toxic exposure."
    },

    "textile": {
        "Recycling type": "Textile reuse",
        "Disposal": "Reuse textiles as rags or donate them before recycling.",
        "Tips": "Keep textiles clean and dry.",
        "Extra": "Textiles can be reused even when worn."
    },

    "general trash": {
        "Recycling type": "General disposal",
        "Disposal": "Dispose of non-recyclable waste in trash bins.",
        "Tips": "Sort recyclables before throwing items away.",
        "Extra": "Proper sorting improves recycling success."
    }
}

analysis = [waste_analyzation_v3, waste_analyzation_v2, waste_analyzation_v1]
top_label = None


def classify_pipeline(image):
    global top_label

    predictions = classify_image(image)
    top_label = None
    top_label = max(predictions, key=predictions.get)  # top-1
    return predictions

def analyze_pipeline():
    global top_label

    if top_label is None:
        return "Please classify an image first."
    
    explanation = explain_recycling(top_label)  #EXCEEDS CPU, WILL GIVE INACCURATE ANSWER. DO NOT RETURN
    
    choice = random.randint(0, 2)
    info = analysis[choice][top_label]  
    result = (
        f"• Recycling type: {info['Recycling type']}\n"
        f"• Disposal: {info['Disposal']}\n"
        f"• Tips: {info['Tips']}\n"
        f"• Extra: {info['Extra']}\n"
    )
    return result


# ------------------ GRADIO UI ------------------
#  CSS
custom_css = """


#main-title {
    text-align: center;
    color: #2E7D32;
    font-size: 34px;
    font-weight: 800;
    margin-bottom: 22px;
}

.gradio-container {
    background: linear-gradient(135deg, #E8F5E9 0%, #F1F8E9 100%);
    font-family: 'Segoe UI', sans-serif;
}

#explainbox textarea {
    background: #ffffff;
    height: 120px; 
    border: 2px solid #A5D6A7;
    border-radius: 12px;
    padding: 12px;
    font-size: 15px;
}

.gr-button.primary {
    background: #43A047 !important;
    color: white !important;
    border-radius: 12px !important;
    padding: 12px 20px !important;
    font-size: 17px !important;
    box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.15);
    transition: 0.2s ease;
}

textarea, .gr-textbox textarea, #explainbox textarea {
    color: #1B5E20 !important;
}

#tips-box li {
    color: #2E7D32 !important;
}

.gr-button.primary:hover {
    background: #2E7D32 !important;
    transform: translateY(-2px);
}


"""

# GRADIO UI
with gr.Blocks() as demo:

    # TITLE
    gr.Markdown("<h1 id='main-title'>♻️ AI Waste Classifier</h1>")

    # INPUT ROW
    with gr.Row():
        with gr.Column(scale=1):
            img_input = gr.Image(
                type="pil",
                label="📸 Upload waste image",
                elem_id="upload-area"
            )

        with gr.Column(scale=1):
            gr.Markdown(
                """
                <div id="tips-box" style="border:2px solid #A5D6A7; padding:16px; border-radius:14px; background:white;">
                    <h3 style="color:#2E7D32;">🌍 Quick recycling tips:</h3>
                    <ul>
                        <li>Organic waste → green bin</li>
                        <li>Plastic, metal waste → recycle</li>
                        <li>Wash your glass waste!</li>
                        <li>Battery, eletrical devices → non-metal container</li>
                    </ul>
                </div>
                """
            )

    # OUTPUTS
    cls_output = gr.Label(
        num_top_classes=3,
        label="🔍 Classifier Prediction (Top 3)"
    )
    analyze_btn = gr.Button("Classify Waste", variant="primary")

    analyze_btn.click(
        classify_pipeline,
        inputs=img_input,
        outputs=cls_output
    )

    explain_output = gr.Textbox(
        label="🧩 Detailed Recycling & Disposal Advice",
        elem_id="explainbox",
        lines=6
    )

    # BUTTON
    analyze_result_btn = gr.Button("Analyze", variant="primary")

    analyze_result_btn.click(
        analyze_pipeline,
        inputs=None,
        outputs=explain_output
    )


demo.launch(css=custom_css)