Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| from tensorflow.keras.models import load_model | |
| from tensorflow.keras.applications.efficientnet import preprocess_input as efficientnet_preprocess | |
| import pandas as pd | |
| import os | |
| from datetime import datetime | |
| # Waste categories and disposal advice | |
| waste_advice = { | |
| "PAPER": { | |
| "category": "Paper", | |
| "advice": "Recyclable: Clean paper, newspapers, magazines, cardboard. Remove staples and plastic windows. Avoid: Wax-coated paper, tissues, heavily soiled paper.", | |
| "bin": "Blue Recycling Bin" | |
| }, | |
| "PLASTIC": { | |
| "category": "Plastic", | |
| "advice": "Check recycling numbers (1-7). Commonly accepted: bottles, containers. Clean containers before recycling. Take plastic bags to store drop-off points.", | |
| "bin": "Yellow Recycling Bin" | |
| }, | |
| "GLASS": { | |
| "category": "Glass", | |
| "advice": "Rinse bottles and jars. Remove metal lids. Sort by colour if required.", | |
| "bin": "Green Recycling Bin" | |
| }, | |
| "METAL": { | |
| "category": "Metal", | |
| "advice": "Aluminium cans are highly recyclable. Remove labels and rinse steel cans. Take larger items to scrapyards.", | |
| "bin": "Red Recycling Bin" | |
| }, | |
| "ORGANIC": { | |
| "category": "Organic", | |
| "advice": "Compost fruit/vegetable scraps, coffee grounds, eggshells. Avoid meat and dairy in home compost.", | |
| "bin": "Brown Compost Bin" | |
| }, | |
| "BATTERIES": { | |
| "category": "Battery", | |
| "advice": "Take to designated collection points at supermarkets or electronics stores. Never put in regular rubbish due to toxic materials.", | |
| "bin": "Special Collection Point" | |
| }, | |
| "E-WASTE": { | |
| "category": "E-Waste", | |
| "advice": "Take to manufacturer take-back programmes or electronics retailers. Wipe personal data before disposal. Contains valuable materials.", | |
| "bin": "E-Waste Collection Centre" | |
| }, | |
| "CLOTHES": { | |
| "category": "Clothes", | |
| "advice": "Donate clean, wearable items. Use textile recycling bins for worn-out clothes. Many retailers offer take-back programmes.", | |
| "bin": "Donation Centre/Textile Bin" | |
| }, | |
| "LIGHT-BULBS": { | |
| "category": "Light Bulbs", | |
| "advice": "Do not place in general recycling. Take CFL and fluorescent bulbs to special collection sites.", | |
| "bin": "Special Disposal Site" | |
| } | |
| } | |
| class_names = ['BATTERIES', 'CLOTHES', 'E-WASTE', 'GLASS', 'LIGHT-BULBS', 'METAL', 'ORGANIC', 'PAPER', 'PLASTIC'] | |
| # Load model | |
| try: | |
| model = load_model("model_efnet.keras") | |
| except OSError: | |
| model = None | |
| print("⚠️ Model file not found. Please ensure 'model_efnet.keras' exists in the directory.") | |
| img_size = (224, 224) | |
| log_file = "prediction_log.csv" | |
| if not os.path.exists(log_file): | |
| pd.DataFrame(columns=["Timestamp", "Prediction", "Confidence", "Feedback"]).to_csv(log_file, index=False) | |
| def predict(image): | |
| if model is None: | |
| return "Model not loaded", "", None, None, {} | |
| img = image.convert("RGB").resize(img_size) | |
| if not isinstance(image, Image.Image): | |
| return "Invalid image", "", None, None, {} | |
| img_array = efficientnet_preprocess(np.array(img)) | |
| img_array = np.expand_dims(img_array, axis=0) | |
| prediction = model.predict(img_array, verbose=0)[0] | |
| predicted_index = np.argmax(prediction) | |
| predicted_class = class_names[predicted_index] | |
| confidence = prediction[predicted_index] | |
| advice_info = waste_advice.get(predicted_class, { | |
| "category": predicted_class, | |
| "advice": "No disposal advice available.", | |
| "bin": "Not specified" | |
| }) | |
| formatted_advice = f""" | |
| <h3><b>Category:</b> {advice_info['category']}</h3> | |
| <p><b>Advice:</b> {advice_info['advice']}</p> | |
| <p><b>Bin:</b> {advice_info['bin']}</p> | |
| """ | |
| result = f"{predicted_class} ({confidence:.2%})" | |
| probabilities = {class_names[i]: float(prediction[i]) for i in range(len(class_names))} | |
| return result, formatted_advice, predicted_class, confidence, probabilities | |
| def save_feedback(predicted_class, confidence, feedback): | |
| timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S') | |
| if confidence is None: | |
| confidence_value = "" | |
| else: | |
| confidence_value = round(confidence, 2) | |
| df = pd.read_csv(log_file) | |
| new_row = pd.DataFrame({ | |
| "Timestamp": [timestamp], | |
| "Prediction": [predicted_class], | |
| "Confidence": [round(confidence, 2)], | |
| "Feedback": [feedback] | |
| }) | |
| df = pd.concat([df, new_row], ignore_index=True) | |
| df.to_csv(log_file, index=False) | |
| return( | |
| gr.update(visible=True), # feedback_ack | |
| gr.update(value="", visible=False), # feedback_box | |
| gr.update(value=None) # thumbs | |
| ) | |
| def handle_feedback(feedback): | |
| if feedback == "👎 Wrong": | |
| return ( | |
| gr.update(visible=True), # Show textbox | |
| gr.update(visible=True), # Show submit button | |
| gr.update(visible=False) # Hide feedback_ack | |
| ) | |
| return ( | |
| gr.update(visible=False), # Hide textbox | |
| gr.update(visible=False), # Hide submit button | |
| gr.update(value="Thank you", visible=True) # Show feedback_ack with message | |
| ) | |
| def clear_all(): | |
| return None, "", "", None, None, {}, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) | |
| def load_images_from_folder(): | |
| folder_path = "sample images" # Your folder name | |
| image_files = [] | |
| # Get all image files from the folder | |
| for filename in os.listdir(folder_path): | |
| if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')): | |
| image_path = os.path.join(folder_path, filename) | |
| image_files.append(image_path) | |
| return image_files | |
| # Gradio UI | |
| with gr.Blocks() as demo: | |
| gr.Markdown(""" | |
| # ♻️ Smart Waste Classifier | |
| Upload a waste image or Use webcam to get a predicted category and disposal advice. | |
| """) | |
| with gr.Row(): | |
| image_input = gr.Image(label="Upload or Use Webcam", sources=["upload", "webcam"], type="pil") | |
| with gr.Column(): | |
| result = gr.Label(label="Prediction") | |
| advice_output = gr.HTML(label="Waste Disposal Advice") | |
| pred_class = gr.State() | |
| pred_conf = gr.State() | |
| probabilities = gr.State() | |
| with gr.Row(): | |
| predict_button = gr.Button("Detect waste") | |
| clear_btn = gr.Button("Clear") | |
| predict_button.click( | |
| fn=predict, | |
| inputs=image_input, | |
| outputs=[result, advice_output, pred_class, pred_conf, probabilities] | |
| ) | |
| gr.Markdown("### How accurate was the prediction?") | |
| with gr.Row(): | |
| thumbs = gr.Radio(["👍 Correct", "👎 Wrong"], label="Prediction Status") | |
| feedback_box = gr.Textbox(lines=2, placeholder="What should the model have predicted?", visible=False, label="Your Feedback") | |
| submit_feedback = gr.Button("Submit Feedback", visible=False) | |
| feedback_ack = gr.Textbox(visible=False, label="✅ Feedback saved. Thank you!", interactive=False) | |
| clear_btn.click( | |
| fn=clear_all, | |
| outputs=[image_input, result, advice_output, pred_class, pred_conf, probabilities, thumbs, feedback_box, submit_feedback, feedback_ack] | |
| ) | |
| thumbs.change( | |
| fn=handle_feedback, | |
| inputs=[thumbs], | |
| outputs=[feedback_box, submit_feedback, feedback_ack] | |
| ) | |
| submit_feedback.click( | |
| fn=save_feedback, | |
| inputs=[pred_class, pred_conf, feedback_box], | |
| outputs=[feedback_ack, feedback_box, thumbs] | |
| ) | |
| # Inside Gradio UI | |
| gr.Markdown("#### ✨Image Gallery: Test with these sample images") | |
| # Display images in a gallery | |
| gallery = gr.Gallery( | |
| value=load_images_from_folder(), | |
| label="Sample Images", | |
| show_label=True, | |
| elem_id="gallery", | |
| columns=6, | |
| rows=1, | |
| height="auto" | |
| ) | |
| demo.launch() | |