import os import time import csv import json import math import base64 import datetime import threading import collections import numpy as np from PIL import Image, ImageDraw, ImageFont, ImageFilter import io import torch import torch.nn.functional as F from transformers import AutoImageProcessor, AutoModelForImageClassification from flask import Flask, request, jsonify, render_template, Response import requests as http_requests # Optional imports — graceful fallback if not installed try: import cv2 CV2_AVAILABLE = True except ImportError: CV2_AVAILABLE = False print("[WARN] cv2 not available — CLAHE disabled. Run: pip install opencv-python") try: from scipy.ndimage import convolve as scipy_convolve SCIPY_AVAILABLE = True except ImportError: SCIPY_AVAILABLE = False print("[WARN] scipy not available — Wiener deconvolution disabled.") app = Flask(__name__) # ============================================================================= # PATHS # ============================================================================= BASE_DIR = os.path.dirname(os.path.abspath(__file__)) LOGS_DIR = os.path.join(BASE_DIR, "logs") CONFIG_PATH = os.path.join(BASE_DIR, "config", "category_mapping.json") ROI_PATH = os.path.join(BASE_DIR, "config", "roi_config.json") BACKGROUND_PATH = os.path.join(BASE_DIR, "config", "background_ref.npy") MEMORY_PATH = os.path.join(BASE_DIR, "config", "correction_memory.json") CSV_LOG_PATH = os.path.join(LOGS_DIR, "classification_log.csv") os.makedirs(LOGS_DIR, exist_ok=True) os.makedirs(os.path.join(BASE_DIR, "config"), exist_ok=True) # ============================================================================= # GLOBAL STATE # ============================================================================= model = None processor = None category_mapping = {} last_processed_image = None device = "cpu" # ── Frame pipeline ────────────────────────────────────────────────────────── latest_raw_frame = None latest_annotated = None frame_lock = threading.Lock() # ── Background reference ────────────────────────────────────────────────────── background_gray = None background_lock = threading.Lock() background_ready = False # ── Detection state ─────────────────────────────────────────────────────────── reference_gray = None reference_lock = threading.Lock() detect_confirm = 0 PIXEL_DIFF_THRESH = 25 CHANGED_PCT = 0.08 CONFIRM_FRAMES = 3 ADAPT_RATE = 0.04 settle_until = 0 SETTLE_DELAY_S = 1.5 # ── Command queue ───────────────────────────────────────────────────────────── command_queue = [] command_lock = threading.Lock() # ── Classification result ───────────────────────────────────────────────────── last_category = None last_category_time = 0 LABEL_SHOW_S = 5.0 cycle_running = False # ── Gemma agent state ───────────────────────────────────────────────────────── last_gemma_result = { "status": "idle", "prediction": None, "timestamp": None, "upcycling_tip": None, "eco_fact": None, "prep_instruction": None, "witty_remark": None, "material_analysis": None, "highlighted_type": None, "highlighted_title": None, "highlighted_text": None } gemma_lock = threading.Lock() gemma_busy = False # ── Memory database ─────────────────────────────────────────────────────────── correction_memory = [] memory_lock = threading.Lock() # ============================================================================= # ROI CONFIG # ============================================================================= DEFAULT_ROI = { "cx": 160, "cy": 20, "angle_deg": 60, "radius_pct": 80, "zoom_factor": 1.0, } roi_config = dict(DEFAULT_ROI) def load_roi(): global roi_config if os.path.exists(ROI_PATH): try: with open(ROI_PATH) as f: roi_config = {**DEFAULT_ROI, **json.load(f)} except Exception: roi_config = dict(DEFAULT_ROI) def save_roi(): with open(ROI_PATH, "w") as f: json.dump(roi_config, f, indent=2) def roi_radius(): return int(240 * roi_config["radius_pct"] / 100) def build_slice_mask(h, w): cx = roi_config["cx"] cy = roi_config["cy"] half_angle = roi_config["angle_deg"] / 2.0 r = roi_radius() ys, xs = np.mgrid[0:h, 0:w] dx = xs - cx dy = ys - cy dist = np.sqrt(dx*dx + dy*dy) angle = np.abs(np.degrees(np.arctan2(dx, dy))) return (dist <= r) & (dist >= 1) & (angle <= half_angle) def roi_bounding_box(): cx = roi_config["cx"] cy = roi_config["cy"] r = roi_radius() half_rad = math.radians(roi_config["angle_deg"] / 2.0) x0 = max(0, int(cx - r * math.sin(half_rad)) - 4) y0 = max(0, cy - 4) x1 = min(320, int(cx + r * math.sin(half_rad)) + 4) y1 = min(240, int(cy + r * math.cos(half_rad)) + 4) return x0, y0, x1, y1 # ============================================================================= # BACKGROUND REFERENCE # ============================================================================= def load_background(): global background_gray, background_ready if os.path.exists(BACKGROUND_PATH): try: background_gray = np.load(BACKGROUND_PATH).astype(np.float32) background_ready = True except Exception as e: print(f"[BG] Failed to load background: {e}") def save_background(jpeg_bytes: bytes): global background_gray, background_ready img = Image.open(io.BytesIO(jpeg_bytes)).convert("L").resize((320, 240)) gray = np.array(img, dtype=np.float32) with background_lock: background_gray = gray background_ready = True np.save(BACKGROUND_PATH, gray) def apply_background_subtraction(img_np: np.ndarray, mask: np.ndarray) -> np.ndarray: with background_lock: bg = background_gray.copy() if background_ready else None if bg is None: return img_np gray = np.mean(img_np, axis=2).astype(np.float32) diff = np.abs(gray - bg) is_background = diff < 18.0 result = img_np.copy() result[~mask] = 0 result[is_background & mask] = 0 return result # ============================================================================= # IMAGE PREPROCESSING PIPELINE # ============================================================================= def apply_clahe(img: Image.Image) -> Image.Image: if not CV2_AVAILABLE: return img try: arr = np.array(img.convert("RGB")) lab = cv2.cvtColor(arr, cv2.COLOR_RGB2LAB) l, a, b = cv2.split(lab) clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(4, 4)) l = clahe.apply(l) lab = cv2.merge([l, a, b]) result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB) return Image.fromarray(result) except Exception: return img def sharpen_defocus(img: Image.Image) -> Image.Image: try: img = img.filter(ImageFilter.UnsharpMask(radius=2.2, percent=175, threshold=2)) img = img.filter(ImageFilter.EDGE_ENHANCE) return img except Exception: return img def _fallback_bfs_component(mask: np.ndarray) -> np.ndarray: coords = np.argwhere(mask) if len(coords) == 0: return mask h, w = mask.shape visited = np.zeros((h, w), dtype=bool) largest_cluster = [] for y, x in coords: if visited[y, x]: continue cluster = [] queue = [(y, x)] visited[y, x] = True while queue: cy, cx = queue.pop(0) cluster.append((cy, cx)) for dy in [-1, 0, 1]: for dx in [-1, 0, 1]: if dy == 0 and dx == 0: continue ny, nx = cy + dy, cx + dx if 0 <= ny < h and 0 <= nx < w: if mask[ny, nx] and not visited[ny, nx]: visited[ny, nx] = True queue.append((ny, nx)) if len(cluster) > len(largest_cluster): largest_cluster = cluster largest_mask = np.zeros((h, w), dtype=bool) for y, x in largest_cluster: largest_mask[y, x] = True return largest_mask def preprocess_for_classification(jpeg_bytes: bytes) -> Image.Image: img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB").resize((320, 240)) img_np = np.array(img) mask = build_slice_mask(240, 320) with background_lock: bg = background_gray.copy() if background_ready else None img_cropped = None if bg is not None: gray = np.mean(img_np, axis=2).astype(np.float32) diff = np.abs(gray - bg) raw_mask = (diff > 18.0) & mask # 1. Apply Median Filter to wipe out isolated discrete spots mask_pil = Image.fromarray((raw_mask * 255).astype(np.uint8)) mask_pil = mask_pil.filter(ImageFilter.MedianFilter(size=5)) filtered_mask = np.array(mask_pil) > 128 # 2. Extract Largest Connected Component try: from scipy.ndimage import label labeled, num_features = label(filtered_mask) if num_features > 0: sizes = np.bincount(labeled.ravel()) sizes[0] = 0 # Ignore background largest_label = np.argmax(sizes) final_mask = (labeled == largest_label) else: final_mask = filtered_mask except Exception: final_mask = _fallback_bfs_component(filtered_mask) object_coords = np.argwhere(final_mask) if len(object_coords) > 0: y_coords = object_coords[:, 0] x_coords = object_coords[:, 1] # Object center cy_obj = np.mean(y_coords) cx_obj = np.mean(x_coords) # Max distance to furthest continuous point dists = np.sqrt((x_coords - cx_obj)**2 + (y_coords - cy_obj)**2) r_max = np.max(dists) # Move boundary back by 20% to capture surrounding area (margin) half_side = max(15, int(r_max * 1.20)) # Construct a perfect square side length side = 2 * half_side side = min(side, 240, 320) # Keep within image bounds half_side = side // 2 ymin = int(cy_obj - half_side) ymax = int(cy_obj + half_side) xmin = int(cx_obj - half_side) xmax = int(cx_obj + half_side) # Shift the window if it goes out of bounds to keep it square if xmin < 0: xmax -= xmin xmin = 0 if xmax > 320: xmin -= (xmax - 320) xmax = 320 if ymin < 0: ymax -= ymin ymin = 0 if ymax > 240: ymin -= (ymax - 240) ymax = 240 img_cropped = img.crop((xmin, ymin, xmax, ymax)) print(f"[AI] Connected Object Square Crop: Center({cx_obj:.1f}, {cy_obj:.1f}), r_max={r_max:.1f}, box=({xmin},{ymin}) to ({xmax},{ymax})") if img_cropped is None: # Fallback to static ROI bounding box x0, y0, x1, y1 = roi_bounding_box() img_cropped = img.crop((x0, y0, x1, y1)) if (x1 > x0 and y1 > y0) else img print("[AI] Fallback to static ROI crop.") # In case the crop is not square (e.g. from static fallback), do center crop to square it w, h = img_cropped.size if w != h: min_dim = min(w, h) left = (w - min_dim) // 2 top = (h - min_dim) // 2 img_cropped = img_cropped.crop((left, top, left + min_dim, top + min_dim)) print(f"[AI] Squared fallback crop: {w}x{h} -> {min_dim}x{min_dim}") # Resize directly and cleanly without aggressive sharpening filters to maintain natural quality img_cropped = img_cropped.resize((224, 224), Image.LANCZOS) return img_cropped # ============================================================================= # OVERLAY DRAWING # ============================================================================= CATEGORY_COLORS = { "BIODEGRADABLE": (0, 200, 0), "NON_BIODEGRADABLE": (220, 0, 0), "OTHER": (220, 200, 0), } SLICE_COLOR = (255, 60, 60) SLICE_ALPHA = 40 def draw_slice_overlay(img: Image.Image) -> Image.Image: cx = roi_config["cx"] cy = roi_config["cy"] half_angle = roi_config["angle_deg"] / 2.0 r = roi_radius() half_rad = math.radians(half_angle) overlay = Image.new("RGBA", img.size, (0, 0, 0, 0)) draw = ImageDraw.Draw(overlay) steps = 60 points = [(cx, cy)] for s in range(steps + 1): a = -half_rad + 2 * half_rad * s / steps ax = int(cx + r * math.sin(a)) ay = int(cy + r * math.cos(a)) points.append((ax, ay)) points.append((cx, cy)) draw.polygon(points, fill=(*SLICE_COLOR, SLICE_ALPHA)) arc_lx = int(cx - r * math.sin(half_rad)) arc_ly = int(cy + r * math.cos(half_rad)) arc_rx = int(cx + r * math.sin(half_rad)) draw.line([(cx, cy), (arc_lx, arc_ly)], fill=(*SLICE_COLOR, 230), width=2) draw.line([(cx, cy), (arc_rx, arc_ly)], fill=(*SLICE_COLOR, 230), width=2) for s in range(steps): a0 = -half_rad + 2 * half_rad * s / steps a1 = -half_rad + 2 * half_rad * (s + 1) / steps p0 = (int(cx + r * math.sin(a0)), int(cy + r * math.cos(a0))) p1 = (int(cx + r * math.sin(a1)), int(cy + r * math.cos(a1))) draw.line([p0, p1], fill=(*SLICE_COLOR, 230), width=2) img = img.convert("RGBA") img = Image.alpha_composite(img, overlay) return img.convert("RGB") def draw_category_tag(img: Image.Image, category: str) -> Image.Image: cx, cy = roi_config["cx"], roi_config["cy"] r = roi_radius() half_rad = math.radians(roi_config["angle_deg"] / 2.0) arc_lx = int(cx - r * math.sin(half_rad)) arc_rx = int(cx + r * math.sin(half_rad)) arc_by = int(cy + r * math.cos(half_rad)) box_x0, box_y0 = max(0, arc_lx - 4), max(0, cy - 20) box_x1, box_y1 = min(319, arc_rx + 4), min(239, arc_by + 4) color = CATEGORY_COLORS.get(category, (180, 180, 0)) draw = ImageDraw.Draw(img) # Bounding box outline only (3px width) for t in range(3): draw.rectangle([(box_x0 + t, box_y0 + t), (box_x1 - t, box_y1 - t)], outline=color) # Top filled strip for label text strip_y1 = min(239, box_y0 + 14) draw.rectangle([(box_x0, box_y0), (box_x1, strip_y1)], fill=color) try: font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 10) except Exception: font = ImageFont.load_default() draw.text((box_x0 + 3, box_y0 + 2), category, fill=(0, 0, 0), font=font) return img def annotate_frame(jpeg_bytes: bytes, category: str = None, show_overlays: bool = False) -> bytes: try: if not show_overlays: return jpeg_bytes img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB") img = draw_slice_overlay(img) if category: img = draw_category_tag(img, category) out = io.BytesIO() img.save(out, format="JPEG", quality=90) return out.getvalue() except Exception: return jpeg_bytes def make_bg_preview(jpeg_bytes: bytes) -> bytes: try: img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB") img_np = np.array(img) mask = build_slice_mask(240, 320) with background_lock: bg = background_gray.copy() if background_ready else None if bg is not None: gray = np.mean(img_np, axis=2).astype(np.float32) diff = np.abs(gray - bg) is_bg_pixel = diff < 18.0 overlay = img_np.copy() bg_mask = is_bg_pixel & mask overlay[bg_mask, 0] = np.clip(overlay[bg_mask, 0].astype(int) * 0.4, 0, 255).astype(np.uint8) overlay[bg_mask, 1] = np.clip(overlay[bg_mask, 1].astype(int) * 0.4 + 80, 0, 255).astype(np.uint8) overlay[bg_mask, 2] = np.clip(overlay[bg_mask, 2].astype(int) * 0.4, 0, 255).astype(np.uint8) outside = ~mask overlay[outside] = (overlay[outside].astype(int) * 0.25).astype(np.uint8) img = Image.fromarray(overlay) else: arr = img_np.copy() arr[~mask] = (arr[~mask].astype(int) * 0.3).astype(np.uint8) img = Image.fromarray(arr) img = draw_slice_overlay(img) out = io.BytesIO() img.save(out, format="JPEG", quality=80) return out.getvalue() except Exception: return jpeg_bytes # ============================================================================= # DETECTION LOOP # ============================================================================= def frame_to_gray(jpeg_bytes: bytes) -> np.ndarray: img = Image.open(io.BytesIO(jpeg_bytes)).convert("L").resize((320, 240)) return np.array(img, dtype=np.float32) def update_reference(gray: np.ndarray): global reference_gray with reference_lock: if reference_gray is None: reference_gray = gray.copy() else: reference_gray = (1 - ADAPT_RATE) * reference_gray + ADAPT_RATE * gray def check_for_waste(gray: np.ndarray) -> bool: global detect_confirm, reference_gray if time.time() < settle_until: return False with reference_lock: if reference_gray is None: return False ref = reference_gray.copy() mask = build_slice_mask(240, 320) diff = np.abs(gray - ref) changed = int(np.sum((diff > PIXEL_DIFF_THRESH) & mask)) total = int(np.sum(mask)) trigger = int(total * CHANGED_PCT) if changed >= trigger: detect_confirm += 1 if detect_confirm >= CONFIRM_FRAMES: detect_confirm = 0 return True else: detect_confirm = 0 update_reference(gray) return False def detection_loop(): global latest_annotated, last_category, last_category_time, cycle_running, settle_until while True: try: time.sleep(0.08) with frame_lock: raw = latest_raw_frame if raw is None: continue gray = frame_to_gray(raw) with reference_lock: have_ref = reference_gray is not None if not have_ref: update_reference(gray) continue # Check if there is active motion (changed pixels exceed threshold) motion_detected = False with reference_lock: if reference_gray is not None: ref = reference_gray.copy() mask = build_slice_mask(240, 320) diff = np.abs(gray - ref) changed = int(np.sum((diff > PIXEL_DIFF_THRESH) & mask)) total = int(np.sum(mask)) motion_detected = (changed >= int(total * CHANGED_PCT)) cat = last_category if last_category and (time.time() - last_category_time) < LABEL_SHOW_S else None # Show overlays only when motion is detected, cycle is running, or recent category is active show_overlays = motion_detected or cycle_running or (cat is not None) annotated = annotate_frame(raw, cat, show_overlays=show_overlays) with frame_lock: latest_annotated = annotated if time.time() < settle_until: continue if check_for_waste(gray): result = classify_image(raw, log_to_csv=False) last_category = result["final_category"] last_category_time = time.time() settle_until = time.time() + 1.25 except Exception: time.sleep(1.0) def wait_for_queue_done(): while True: time.sleep(0.1) with command_lock: empty = len(command_queue) == 0 if empty: global cycle_running, reference_gray cycle_running = False with reference_lock: reference_gray = None break # ============================================================================= # GEMMA AGENT # ============================================================================= OLLAMA_URL = "http://localhost:11434/api/generate" OLLAMA_MODEL = "gemma4" STATIC_FALLBACKS = { "BIODEGRADABLE": { "upcycling_tip": "Compost this to nourish your garden soil or house plants.", "eco_fact": "Organic matter decomposes quickly, but in landfills without oxygen, it produces harmful methane gas.", "prep_instruction": "Remove any plastic packaging, stickers, or metal twist ties before composting.", "witty_remark": "Yum! Earthworms are going to love this delicious biodegradable snack.", "material_analysis": "Organic biodegradable waste (biological matter, cardboard, or clean paper)." }, "NON_BIODEGRADABLE": { "upcycling_tip": "Clean and reuse plastic jars for storage, or turn metal cans into pencil holders.", "eco_fact": "A plastic bottle can take up to 450 years to decompose in a landfill.", "prep_instruction": "Empty all liquids, rinse out food residue, and flatten bottles or cans to save bin space.", "witty_remark": "Beep boop! This non-biodegradable item is going to stick around longer than my server uptime.", "material_analysis": "Recyclable or non-organic materials (plastics, metals, or glass)." }, "OTHER": { "upcycling_tip": "Repurpose fabrics or old shoes into cleaning rags or workshop items.", "eco_fact": "Landfill space is finite, and mixed material items are the hardest to sort and recycle.", "prep_instruction": "Separate recyclable components from trash if possible. Do not put heavily contaminated items in recycling.", "witty_remark": "Sorting is hard, but together we make a great team! Let's handle this item responsibly.", "material_analysis": "General waste or mixed composition items (clothes, shoes, trash)." } } def _gemma_available() -> bool: try: return http_requests.get("http://localhost:11434/api/tags", timeout=2).status_code == 200 except Exception: return False def _image_to_b64(img: Image.Image) -> str: buf = io.BytesIO() img.save(buf, format="PNG") return base64.b64encode(buf.getvalue()).decode() def text_to_robot_tones(text: str) -> str: words = text.split() tones = [] import random rng = random.Random(hash(text)) for word in words: syllables = max(1, len(word) // 3) for _ in range(syllables): freq = rng.randint(600, 1500) duration = rng.randint(45, 95) tones.append(f"{freq},{duration}") return "|".join(tones[:12]) def _run_gemma(img_crop: Image.Image, hf_result: dict) -> dict: cat = hf_result.get("final_category", "OTHER") fallback = STATIC_FALLBACKS.get(cat, STATIC_FALLBACKS["OTHER"]) if not _gemma_available(): res = { "status": "done", "prediction": cat, "timestamp": datetime.datetime.now().isoformat(), **fallback } return res try: img_b64 = _image_to_b64(img_crop) prompt = f""" Analyze this waste item. The primary image classifier has classified this item as: '{hf_result.get("detailed_label")}' (Category: {cat}). You must respond ONLY with a single JSON object containing these exact keys: - "upcycling_tip": A creative upcycling/reuse idea for this specific item. - "eco_fact": An interesting environmental or decomposition fact about this material. - "prep_instruction": How to prepare this item (e.g. rinse, remove caps, flatten, check for grease). - "witty_remark": A short, witty, or encouraging comment in the persona of a friendly eco-conscious robot assistant. - "material_analysis": A detailed breakdown of materials or recycle codes (e.g. PET 1, paper, aluminum). Ensure the response is valid JSON and contains NO other text, no markdown block wrappers (like ```json), and no backticks. """ payload = {"model": OLLAMA_MODEL, "prompt": prompt, "images": [img_b64], "stream": False} r = http_requests.post(OLLAMA_URL, json=payload, timeout=12) raw = r.json().get("response", "").strip() import re match = re.search(r'\{.*\}', raw, re.DOTALL) if match: json_data = json.loads(match.group(0)) else: json_data = json.loads(raw) return { "status": "done", "prediction": cat, "timestamp": datetime.datetime.now().isoformat(), "upcycling_tip": json_data.get("upcycling_tip") or fallback["upcycling_tip"], "eco_fact": json_data.get("eco_fact") or fallback["eco_fact"], "prep_instruction": json_data.get("prep_instruction") or fallback["prep_instruction"], "witty_remark": json_data.get("witty_remark") or fallback["witty_remark"], "material_analysis": json_data.get("material_analysis") or fallback["material_analysis"] } except Exception as e: print(f"[Gemma Error] {e}") return { "status": "done", "prediction": cat, "timestamp": datetime.datetime.now().isoformat(), **fallback } def gemma_analyze_async(img_crop: Image.Image, hf_result: dict, image_name: str): global gemma_busy, last_gemma_result cat = hf_result.get("final_category", "OTHER") with gemma_lock: last_gemma_result = { "status": "analyzing", "prediction": cat, "timestamp": datetime.datetime.now().isoformat(), "upcycling_tip": "...", "eco_fact": "...", "prep_instruction": "...", "witty_remark": "...", "material_analysis": "...", "highlighted_type": "...", "highlighted_title": "...", "highlighted_text": "..." } def _worker(): global gemma_busy, last_gemma_result res = _run_gemma(img_crop, hf_result) import random insights = [ ("upcycling", "💡 Upcycling Tip", res.get("upcycling_tip")), ("fact", "🌍 Eco-Fact", res.get("eco_fact")), ("prep", "⚙️ Prep Info", res.get("prep_instruction")), ("material", "🔍 Material Analysis", res.get("material_analysis")), ("witty", "🤖 Witty Remark", res.get("witty_remark")), ] valid_insights = [i for i in insights if i[2] and i[2] != "..."] if valid_insights: chosen_type, chosen_title, chosen_text = random.choice(valid_insights) else: chosen_type, chosen_title, chosen_text = ("witty", "🤖 Witty Remark", res.get("witty_remark") or "Sorting complete!") res["highlighted_type"] = chosen_type res["highlighted_title"] = chosen_title res["highlighted_text"] = chosen_text with gemma_lock: last_gemma_result = res tones = text_to_robot_tones(chosen_text) with command_lock: command_queue.append({"cmd": "ROBO_TALK", "param": tones}) gemma_busy = False gemma_busy = True threading.Thread(target=_worker, daemon=True).start() # ============================================================================= # MEMORY DATABASE # ============================================================================= def load_memory(): global correction_memory if os.path.exists(MEMORY_PATH): try: with open(MEMORY_PATH) as f: correction_memory = json.load(f) except Exception: correction_memory = [] def save_correction(entry: dict): with memory_lock: correction_memory.append(entry) data = list(correction_memory) try: with open(MEMORY_PATH, "w") as f: json.dump(data, f, indent=2) except Exception: pass # ============================================================================= # AI CLASSIFICATION # ============================================================================= def get_waste_attributes(label: str, category: str) -> dict: label_lower = label.lower() # Recyclable: plastic, cardboard, paper, metal, glass, cans, bottles recyclable_keywords = ["plastic", "cardboard", "paper", "metal", "glass", "can", "bottle", "scale"] is_recyclable = any(x in label_lower for x in recyclable_keywords) # Exclude non-recyclable items if label_lower in ["trash", "battery", "batteries", "electronics", "biological", "clothes", "shoes", "eraser"]: is_recyclable = False # Moisture: biological / organic are wet wet_keywords = ["biological", "organic", "food", "wet"] is_wet = any(x in label_lower for x in wet_keywords) return { "recyclable": "RECYCLABLE" if is_recyclable else "NON_RECYCLABLE", "moisture": "WET" if is_wet else "DRY" } def classify_image(jpeg_bytes: bytes, log_to_csv: bool = True, bypass_preprocess: bool = False) -> dict: global last_processed_image try: if bypass_preprocess: img = Image.open(io.BytesIO(jpeg_bytes)).convert("RGB") # Center crop to a perfect square to prevent squeezing/stretching w, h = img.size if w != h: min_dim = min(w, h) left = (w - min_dim) // 2 top = (h - min_dim) // 2 img = img.crop((left, top, left + min_dim, top + min_dim)) print(f"[AI] Bypassing classical CV crop for browser frame; center-cropped to square ({w}x{h} -> {img.size[0]}x{img.size[1]}).") else: img = preprocess_for_classification(jpeg_bytes) # Save preprocessed image for debug/viewing try: img.save(os.path.join(LOGS_DIR, "cropped_object.jpg"), format="JPEG", quality=90) except Exception: pass inputs = processor(images=img, return_tensors="pt") inputs = {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): outputs = model(**inputs) probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0] max_prob, max_idx = torch.max(probs, dim=-1) prob = max_prob.item() * 100 raw_label = model.config.id2label[max_idx.item()] raw_label_lower = raw_label.lower() # Translate ResNet French labels to standard English mapping keywords french_to_english = { "batterie": "battery", "carton": "cardboard", "metal": "metal", "organique": "biological", "papier": "paper", "plastique": "plastic", "verre": "glass", "vetements": "clothes" } label = french_to_english.get(raw_label_lower, raw_label_lower) label_lower = label.lower() mapping_lower = {k.lower(): v for k, v in category_mapping.items()} if prob < 75.0 or label_lower not in mapping_lower: category = "OTHER" status = "low_confidence" if prob < 75.0 else "unmapped_label" else: category = mapping_lower[label_lower] status = "mapped" cid = f"waste_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" last_processed_image = cid attributes = get_waste_attributes(label, category) result = { "image_name": cid, "detailed_label": label, "model_confidence": round(prob, 2), "final_category": category, "mapping_status": status, "recyclable": attributes["recyclable"], "moisture": attributes["moisture"] } if log_to_csv: log_prediction_to_csv(datetime.datetime.now().isoformat(), cid, label, prob, category, status, "pending") gemma_analyze_async(img, result, cid) return result except Exception as e: print(f"[AI] Error: {e}") return { "image_name": "error", "detailed_label": "unknown", "model_confidence": 0.0, "final_category": "OTHER", "mapping_status": "error", } # ============================================================================= # CSV LOGGING # ============================================================================= def log_prediction_to_csv(timestamp, image_name, detailed_label, confidence, final_category, mapping_status, confirmation_status): file_exists = os.path.exists(CSV_LOG_PATH) os.makedirs(os.path.dirname(CSV_LOG_PATH), exist_ok=True) with open(CSV_LOG_PATH, mode="a", newline="", encoding="utf-8") as f: writer = csv.writer(f) if not file_exists: writer.writerow(["timestamp","image_name","detailed_label", "model_confidence","final_category","mapping_status","confirmation_status"]) writer.writerow([timestamp, image_name, detailed_label, f"{confidence:.2f}", final_category, mapping_status, confirmation_status]) def update_csv_confirmation(image_name, new_status="confirmed"): if not os.path.exists(CSV_LOG_PATH): return False rows = [] with open(CSV_LOG_PATH, mode="r", newline="", encoding="utf-8") as f: reader = csv.reader(f) for row in reader: rows.append(row) for row in rows: if len(row) > 1 and row[1] == image_name: row[6] = new_status with open(CSV_LOG_PATH, mode="w", newline="", encoding="utf-8") as f: csv.writer(f).writerows(rows) return True # ============================================================================= # COMMAND QUEUE # ============================================================================= def enqueue_sort_cycle(category: str): esp_param = "OTHER" if category == "BIODEGRADABLE": esp_param = "BIO" elif category == "NON_BIODEGRADABLE": esp_param = "NONBIO" with command_lock: command_queue.clear() command_queue.extend([ {"cmd": "SET_LED", "param": category}, {"cmd": "BEEP"}, {"cmd": "SERVO_A", "param": esp_param}, {"cmd": "SERVO_B", "param": "OPEN"}, {"cmd": "DELAY", "param": "750"}, {"cmd": "SERVO_B", "param": "CLOSE"}, {"cmd": "CLEAR_LEDS"}, {"cmd": "SERVO_A", "param": "HOME"}, {"cmd": "BEEP"} ]) # ============================================================================= # STARTUP # ============================================================================= def load_resources(): global model, processor, category_mapping, device with open(CONFIG_PATH) as f: category_mapping = json.load(f) model_path = "dan-lara/Garbage-Classifier-Resnet-50-Finetuning" print(f"[AI] Loading model from Hugging Face Hub: {model_path}") device = "cuda" if torch.cuda.is_available() else "cpu" print(f"[AI] Running on device: {device}") # Workaround: Model repo has pytorch_model.pth instead of pytorch_model.bin from huggingface_hub import snapshot_download import shutil try: local_dir = snapshot_download(model_path) pth_path = os.path.join(local_dir, "pytorch_model.pth") bin_path = os.path.join(local_dir, "pytorch_model.bin") if os.path.exists(pth_path) and not os.path.exists(bin_path): print(f"[AI] Copying pytorch_model.pth to pytorch_model.bin: {bin_path}") shutil.copy(pth_path, bin_path) load_path = local_dir except Exception as e: print(f"[AI] Warning: Failed to download snapshot. Trying direct load fallback. Error: {e}") load_path = model_path processor = AutoImageProcessor.from_pretrained(load_path) model = AutoModelForImageClassification.from_pretrained(load_path).to(device) load_roi() load_background() load_memory() # Initialize empty CSV file with header if it doesn't exist if not os.path.exists(CSV_LOG_PATH): os.makedirs(os.path.dirname(CSV_LOG_PATH), exist_ok=True) with open(CSV_LOG_PATH, mode="w", newline="", encoding="utf-8") as f: csv.writer(f).writerow(["timestamp","image_name","detailed_label", "model_confidence","final_category","mapping_status","confirmation_status"]) threading.Thread(target=detection_loop, daemon=True).start() # ============================================================================= # FLASK ROUTES # ============================================================================= @app.route("/", methods=["GET"]) def index(): return render_template("index.html") def apply_zoom(jpeg_bytes: bytes) -> bytes: zoom = roi_config.get("zoom_factor", 1.0) if zoom <= 1.0: return jpeg_bytes try: img = Image.open(io.BytesIO(jpeg_bytes)) w, h = img.size crop_w = w / zoom crop_h = h / zoom x0 = (w - crop_w) / 2 y0 = (h - crop_h) / 2 x1 = x0 + crop_w y1 = y0 + crop_h img_cropped = img.crop((x0, y0, x1, y1)).resize((w, h), Image.LANCZOS) out = io.BytesIO() img_cropped.save(out, format="JPEG", quality=85) return out.getvalue() except Exception as e: print(f"[ZOOM] Preprocessing failed: {e}") return jpeg_bytes @app.route("/api/upload_frame", methods=["POST"]) def upload_frame(): global latest_raw_frame if not request.data: return jsonify({"error": "No data"}), 400 zoomed = apply_zoom(request.data) with frame_lock: latest_raw_frame = zoomed return "", 200 @app.route("/api/live_feed", methods=["GET"]) def live_feed(): with frame_lock: frame = latest_annotated or latest_raw_frame return Response(frame, mimetype="image/jpeg") if frame else ("", 404) @app.route("/api/command", methods=["GET"]) def get_command(): with command_lock: cmd = command_queue[0] if command_queue else {"cmd": "NONE"} return jsonify(cmd) @app.route("/api/command_done", methods=["POST"]) def command_done(): with command_lock: if command_queue: command_queue.pop(0) if not command_queue: threading.Thread(target=wait_for_queue_done, daemon=True).start() return "", 200 @app.route("/api/user_connected", methods=["POST"]) def user_connected(): with command_lock: command_queue.append({"cmd": "PLAY_NOKIA", "param": ""}) return jsonify({"status": "queued"}) @app.route("/api/roi", methods=["GET", "POST"]) def manage_roi(): if request.method == "POST": data = request.get_json() for key in ("cx", "cy", "angle_deg", "radius_pct", "zoom_factor"): if key in data: roi_config[key] = float(data[key]) save_roi() return jsonify(roi_config) @app.route("/api/roi/reset", methods=["POST"]) def reset_roi(): global roi_config roi_config = dict(DEFAULT_ROI) save_roi() return jsonify({"status": "ok", "roi": roi_config}), 200 @app.route("/api/capture_background", methods=["POST"]) def capture_background(): with frame_lock: raw = latest_raw_frame if not raw: return jsonify({"error": "No frame"}), 400 save_background(raw) return jsonify({"status": "ok"}) @app.route("/api/background_preview", methods=["GET"]) def background_preview(): with frame_lock: raw = latest_raw_frame return Response(make_bg_preview(raw), mimetype="image/jpeg") if raw else ("", 404) @app.route("/api/debug_crop", methods=["GET"]) def debug_crop(): crop_path = os.path.join(LOGS_DIR, "cropped_object.jpg") if os.path.exists(crop_path): with open(crop_path, "rb") as f: data = f.read() return Response(data, mimetype="image/jpeg") return "No cropped image available yet.", 404 @app.route("/classify", methods=["POST"]) def classify_endpoint(): global latest_raw_frame, last_category, last_category_time, settle_until, cycle_running # Check trigger_hardware parameter trigger_hardware = request.args.get("trigger_hardware", "true").lower() == "true" if trigger_hardware and cycle_running: return jsonify({"error": "Cycle already in progress"}), 409 image_data = request.data or latest_raw_frame if not image_data: return jsonify({"error": "No image data"}), 400 if trigger_hardware: cycle_running = True is_browser = bool(request.data) result = classify_image(image_data, log_to_csv=True, bypass_preprocess=is_browser) category = result.get("final_category") if category: if trigger_hardware: enqueue_sort_cycle(category) last_category = category last_category_time = time.time() settle_until = time.time() + SETTLE_DELAY_S + 4.0 else: last_category = category last_category_time = time.time() else: if trigger_hardware: cycle_running = False return jsonify(result), 200 @app.route("/api/trigger_sort", methods=["POST"]) def trigger_sort(): global latest_raw_frame, last_category, last_category_time, settle_until, cycle_running if cycle_running: return jsonify({"status": "ignored", "reason": "cycle_active"}), 200 with frame_lock: img_data = latest_raw_frame if not img_data: return jsonify({"status": "ignored", "reason": "no_frame"}), 200 cycle_running = True print("[TRIGGER] Hardware triggered sorting cycle.") result = classify_image(img_data, log_to_csv=True) category = result.get("final_category", "OTHER") enqueue_sort_cycle(category) last_category = category last_category_time = time.time() settle_until = time.time() + SETTLE_DELAY_S + 4.0 return jsonify({"status": "triggered", "category": category, "result": result}), 200 @app.route("/api/last_gemma", methods=["GET"]) def get_last_gemma(): with gemma_lock: res = dict(last_gemma_result) return jsonify(res), 200 @app.route("/api/feedback", methods=["POST"]) def feedback(): data = request.get_json(silent=True) or {} with gemma_lock: gr = dict(last_gemma_result) entry = { "timestamp": datetime.datetime.now().isoformat(), "image_name": data.get("image_name", gr.get("image_name", "")), "upcycling_tip": gr.get("upcycling_tip", ""), "eco_fact": gr.get("eco_fact", ""), "prep_instruction": gr.get("prep_instruction", ""), "witty_remark": gr.get("witty_remark", ""), "material_analysis": gr.get("material_analysis", ""), "hf_prediction": gr.get("prediction", ""), "user_confirmed": data.get("confirmed_category", ""), "was_correct": data.get("was_correct", True), } save_correction(entry) # Also update CSV update_csv_confirmation(entry["image_name"], "confirmed") return jsonify({"status": "saved"}), 200 # ── System state ────────────────────────────────────────────────────────────── @app.route("/api/is_active", methods=["GET"]) def is_active(): return jsonify({"active": True}), 200 @app.route("/api/handshake", methods=["GET"]) def handshake(): global cycle_running, reference_gray, command_queue with command_lock: command_queue.clear() with reference_lock: reference_gray = None cycle_running = False print("[HANDSHAKE] State synchronized.") return jsonify({"status": "synchronized"}), 200 @app.route("/confirm", methods=["POST"]) def confirm_classification(): global last_processed_image data = request.get_json(silent=True) or {} image_name = data.get("image_name") or last_processed_image if not image_name: return jsonify({"error": "No transaction to confirm"}), 400 success = update_csv_confirmation(image_name) if success: return jsonify({"status": "success", "confirmed_image": image_name}), 200 return jsonify({"error": "Not found"}), 404 # ── Categories ──────────────────────────────────────────────────────────────── @app.route("/api/categories", methods=["GET", "POST"]) def manage_categories(): global category_mapping if request.method == "POST": try: new_mapping = request.get_json() if not new_mapping or not isinstance(new_mapping, dict): return jsonify({"error": "Invalid mapping"}), 400 with open(CONFIG_PATH, "w") as f: json.dump(new_mapping, f, indent=2) category_mapping = new_mapping return jsonify({"status": "success"}), 200 except Exception as e: return jsonify({"error": str(e)}), 500 return jsonify(category_mapping), 200 # ── Logs ────────────────────────────────────────────────────────────────────── @app.route("/api/logs", methods=["GET"]) def get_logs(): if not os.path.exists(CSV_LOG_PATH): return jsonify([]), 200 try: with open(CSV_LOG_PATH, mode="r", newline="", encoding="utf-8") as f: logs = list(csv.DictReader(f)) return jsonify(logs[::-1][:20]), 200 except Exception as e: return jsonify({"error": str(e)}), 500 # ── Background Status ───────────────────────────────────────────────────────── @app.route("/api/background_status", methods=["GET"]) def background_status(): return jsonify({"ready": background_ready}), 200 # ============================================================================= if __name__ == "__main__": load_resources() port = int(os.environ.get("PORT", 7860)) app.run(host="0.0.0.0", port=port, debug=False, threaded=True)