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: \n" "• Disposal: \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("

♻️ AI Waste Classifier

") # 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( """

🌍 Quick recycling tips:

  • Organic waste → green bin
  • Plastic, metal waste → recycle
  • Wash your glass waste!
  • Battery, eletrical devices → non-metal container
""" ) # 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)