import gradio as gr import numpy as np from PIL import Image import os import json import re import tensorflow as tf # -------------------------- # MODEL LOADING # -------------------------- MODEL_FILES = [ "best_waste_classification_model.h5", "cnn_best_model.h5", "mobilenet_enhanced_complete_model.h5" ] model = None for mf in MODEL_FILES: if os.path.exists(mf): try: model = tf.keras.models.load_model(mf) print(f"Loaded model: {mf}") break except Exception as e: print("Failed loading model", mf, e) # -------------------------- # LOAD LABELS # -------------------------- labels = [] if os.path.exists("labels.txt"): with open("labels.txt", "r", encoding="utf-8") as f: labels = [line.strip() for line in f.readlines() if line.strip()] print("Loaded labels (count={}): {}".format(len(labels), labels[:50])) # -------------------------- # LOAD TIPS JSON # -------------------------- tips = {} if os.path.exists("recycling_tips.json"): try: with open("recycling_tips.json", "r", encoding="utf-8") as f: tips = json.load(f) except Exception as e: print("Error loading recycling_tips.json:", e) print("Loaded tips keys (count={}): {}".format(len(tips), list(tips.keys())[:50])) TARGET_SIZE = (224, 224) # -------------------------- # Normalization & lookup helpers # -------------------------- def normalize_key(s: str) -> str: """Lowercase, replace underscores/dashes with spaces, remove non-alnum except spaces, collapse spaces.""" if s is None: return "" s = str(s).strip().lower() s = re.sub(r"[_\-]+", " ", s) s = re.sub(r"[^a-z0-9 ]+", "", s) s = re.sub(r"\s+", " ", s).strip() return s # Build lookup maps tips_exact = dict(tips) # original tips_lower = {k.lower(): v for k, v in tips.items()} tips_norm = {normalize_key(k): v for k, v in tips.items()} tips_norm_keys = list(tips_norm.keys()) def format_raw_tip_to_markdown(raw_tip): """Convert the raw JSON value into a Markdown string (with bullets and headers).""" if isinstance(raw_tip, dict): parts = [] # Tips as bullet list if "tips" in raw_tip and isinstance(raw_tip["tips"], list): parts.append("### Tips") for t in raw_tip["tips"]: parts.append(f"- {t}") # Preparation if "preparation" in raw_tip: parts.append("") parts.append("**Preparation:**") parts.append(raw_tip["preparation"]) # Recyclability if "recyclability" in raw_tip: parts.append("") parts.append("**Recyclability:**") parts.append(raw_tip["recyclability"]) # Join with double newlines so Markdown renders cleanly return "\n\n".join([p for p in parts if p is not None]) else: # Plain string -> return as paragraph return str(raw_tip) # -------------------------- # Preprocess # -------------------------- def preprocess_image(img: Image.Image): img = img.convert("RGB") img = img.resize(TARGET_SIZE) arr = np.array(img) / 255.0 arr = np.expand_dims(arr, axis=0).astype(np.float32) return arr # -------------------------- # Main predict with robust lookup # -------------------------- def predict(image): # Always return two outputs: prediction_text, tip_markdown if image is None: return "No image provided", "" if model is None: return "Model not found or failed to load", "" try: x = preprocess_image(image) preds = model.predict(x) preds = np.asarray(preds).squeeze() if preds.ndim == 0: idx = 0 confidence = float(preds) else: idx = int(np.argmax(preds)) confidence = float(np.max(preds)) except Exception as e: print("Error during prediction:", e) return f"Error: {e}", "" raw_label = labels[idx] if (0 <= idx < len(labels)) else f"class_{idx}" norm_label = normalize_key(raw_label) raw_tip = None source = None # 1) exact original key if raw_label in tips_exact: raw_tip = tips_exact[raw_label]; source = "exact" # 2) exact lowercase if raw_tip is None and raw_label.lower() in tips_lower: raw_tip = tips_lower[raw_label.lower()]; source = "lower" # 3) normalized match if raw_tip is None and norm_label in tips_norm: raw_tip = tips_norm[norm_label]; source = "normalized" # 4) token overlap fuzzy match if raw_tip is None: tokens = set(norm_label.split()) best_match = None best_score = 0 for tk in tips_norm_keys: tk_tokens = set(tk.split()) common = len(tokens & tk_tokens) score = common / (1 + len(tk_tokens)) if score > best_score and common > 0: best_score = score best_match = tk if best_match: raw_tip = tips_norm[best_match]; source = f"token_match:{best_match}" # 5) substring fallback if raw_tip is None: for k in tips_exact.keys(): if k.lower() in raw_label.lower() or raw_label.lower() in k.lower(): raw_tip = tips_exact[k]; source = f"substring:{k}" break # Format tip as Markdown (for large, wrap-friendly display) if raw_tip is None: tip_markdown = "No recycling tip available for this detected label. Please check local recycling rules." else: tip_markdown = format_raw_tip_to_markdown(raw_tip) prediction_text = f"{raw_label} ({confidence*100:.2f}%)" # Debugging log line (useful if something still fails) print("PREDICTION DEBUG:", { "idx": idx, "raw_label": raw_label, "normalized_label": norm_label, "confidence": confidence, "tip_source": source, "tip_found": raw_tip is not None, "available_tip_keys_sample": list(tips_exact.keys())[:50] }) return prediction_text, tip_markdown # -------------------------- # UI - use Markdown for tips so it expands cleanly # -------------------------- with gr.Blocks() as demo: gr.Markdown("# ♻️ Automated Waste Classifier — Tips (expanded view)") with gr.Row(): img = gr.Image(type="pil", label="Upload Image") out = gr.Textbox(label="Prediction", interactive=False, lines=1) # Use Markdown component for the tip so it renders nicely and expands tip_md = gr.Markdown("**Recycling tip will appear here**", label="Recycling Tip") btn = gr.Button("Classify") # Make sure outputs match return values: [prediction_text, tip_markdown] btn.click(fn=predict, inputs=img, outputs=[out, tip_md]) if __name__ == "__main__": demo.launch()