# python .\src\app.py # ------------------------------ # Recycle Material Classifier App # ------------------------------ # This script: # 1. Loads a trained ResNet-18 model # 2. Lets user upload an image or use a live IP camera # 3. Classifies the item (paper/plastic/metal) # 4. Shows Grad-CAM heatmaps for explainability # 5. Displays classification history # ------------------------------ import json, torch from pathlib import Path from PIL import Image from torchvision import transforms import gradio as gr from model import build_model import cv2 import threading import time from explain import generate_gradcam # ---- GLOBAL FLAG (used to stop live feed thread) --- stop_flag = False # ---- MODEL FILE PATHS ---- WEIGHTS = Path("models/resnet18_best.pt") LABELS = Path("models/labels.json") # ---- SELECT DEVICE (GPU if available) ---- device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ---- LOAD LABELS ---- with open(LABELS) as f: idx2name = {int(k): v for k, v in json.load(f).items()} class_names = [idx2name[i] for i in sorted(idx2name.keys())] # ---- LOAD MODEL ---- model = build_model(num_classes=len(class_names), freeze_backbone=False, device=device) state = torch.load(WEIGHTS, map_location=device) model.load_state_dict(state) model.eval() # ---- IMAGE TRANSFORMATIONS ---- # Resize -> Tensor -> Normalize (same as training) tfm = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]), ]) # ---- PREDICTION FUNCTION ---- def predict(img: Image.Image): # Generate Grad-CAM heatmaps (explainable visualization) overlay, heatmap, pred_label, conf = generate_gradcam(img, model, device, class_names) # Compute probability scores for all classes with torch.no_grad(): x = tfm(img.convert("RGB")).unsqueeze(0).to(device) probs = torch.softmax(model(x), dim=1).squeeze(0).cpu().tolist() scores = {cls: float(probs[i]) for i, cls in enumerate(class_names)} top = max(scores, key=scores.get) return [img, overlay, heatmap], pred_label, conf, scores # ---- HISTORY SETTINGS ---- MAX_HISTORY = 12 # show up to 12 previous uploads def classify_and_update(img, history_state): if img is None: return [], "N/A", "N/A", {}, history_state # Run classification gallery_imgs, pred_label, conf, all_scores = predict(img) # Update history (keep last 12 images) history_state.append(img) history_state = history_state[-MAX_HISTORY:] # Pad empty slots padded = history_state + [None]*(MAX_HISTORY - len(history_state)) return gallery_imgs, pred_label, f"{round(conf*100)}%", all_scores, *padded, history_state # ---- HISTORY CLICK EVENT ---- def on_history_select(evt: gr.SelectData, history_state): return history_state[evt.index] # ---- history click ---- def on_history_click(idx, history_state): if idx < len(history_state): return history_state[idx] return None # ---- IP CAMERA SETUP ---- # Replace the IP with your phone’s IP Webcam URL # ip_url = "http://10.132.39.1:8080/video" # replace with your phone's IP ip_url = "http://192.168.1.4:8080/video" # ip_url = "http://10.132.39.1:8080/video" # Variables for motion detection # cap = None # prev_gray = None # motion_active = False # recent_preds = [] def start_live_feed(): global stop_flag stop_flag = False def run(): while not stop_flag: outputs = live_ipcam_generator() # Returns (json_dict, label_dict) json_out_live.update(outputs[0]) label_out_live.update(outputs[1]) time.sleep(0.1) threading.Thread(target=run, daemon=True).start() def stop_live_feed(): global stop_flag stop_flag = True cap = None prev_gray = None motion_active = False recent_preds = [] #ip_url = "http://10.132.39.1:8080/video" #ip_url = "http://192.168.1.6:8080/video" def live_ipcam_generator(): """ Generator that yields only frames with motion detected. Skips all frames without meaningful motion. """ global cap, prev_gray motion_threshold = 100 # How sensitive to motion cooldown_sec = 0.5 # Avoid multiple detections per second last_trigger_time = 0 while True: # Initialize camera if not already if cap is None or not cap.isOpened(): try: cap = cv2.VideoCapture(ip_url) time.sleep(1) ret, prev = cap.read() if not ret or prev is None: prev_gray = None raise ValueError("No frame received") prev_gray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY) except Exception: # If camera fails, send a blank image + "offline" message dummy_img = Image.new("RGB", (224, 224), (0, 0, 0)) yield {"label": "Camera offline", "conf": 0}, {}, [dummy_img], {"motion_level": 0} time.sleep(1) continue # Read frame ret, frame = cap.read() if not ret or frame is None: cap.release() cap = None dummy_img = Image.new("RGB", (224, 224), (0, 0, 0)) yield {"label": "Camera disconnected", "conf": 0}, {}, [dummy_img], {"motion_level": 0} time.sleep(1) continue # Convert to grayscale for motion detection gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if prev_gray is not None: diff = cv2.absdiff(prev_gray, gray) motion_level = cv2.countNonZero(cv2.threshold(diff, 25, 255, cv2.THRESH_BINARY)[1]) else: motion_level = 0 prev_gray = gray # Only process frames with motion above threshold if motion_level > motion_threshold: current_time = time.time() if current_time - last_trigger_time >= cooldown_sec: img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) img = img.resize((840, 480)) pred_images, pred_label, conf, scores = predict(img) pred_json = {"label": pred_label, "conf": round(conf * 100, 2)} motion_info = {"motion_level": motion_level} last_trigger_time = current_time yield pred_json, scores, pred_images, motion_info else: # Skip frame due to cooldown continue else: # Skip frames without motion continue # tiny sleep to avoid hogging CPU time.sleep(0.01) # ---- SIMPLE CSS (hide Gradio footer) ---- css = """ footer, #footer, .footer, [data-testid="branding"] {display:none !important;} a[href*="gradio.app"] {display:none !important;} """ # ---- GRADIO APP LAYOUT ---- with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo: gr.Markdown("

♻️ Recycle Material Classifier

") gr.Markdown("Upload a photo of a recyclable item to classify it as **paper**, **plastic**, or **metal**.") with gr.Tabs(): # --- Upload Image --- # with gr.TabItem("Upload Image"): # img_input = gr.Image(type="pil", label=" Upload an image") # predict_btn = gr.Button("Predict") # # Side-by-side gallery + bar chart # gallery_out = gr.Gallery(label="Original & Grad-CAM", columns=2, height=300) # label_out = gr.Label(num_top_classes=3, label="Top-3 probabilities") # predict_btn.click(predict, inputs=img_input, outputs=[gallery_out, label_out]) # ========== TAB 1: UPLOAD IMAGE ========== with gr.TabItem("Upload Image"): with gr.Row(variant="panel"): # --- Input Column --- with gr.Column(scale=1): image_input = gr.Image( type="pil", label="Upload Image", height=350 ) # Load initial history history_state = gr.State([]) with gr.Row(): history_slots = [ gr.Image(type="pil", interactive=False, height=120, width=120, label=f"#{i+1}") for i in range(MAX_HISTORY) ] # Add select and click events to each history slot for i, slot in enumerate(history_slots): slot.select( fn=lambda h, i=i: on_history_click(i, h), inputs=history_state, outputs=image_input ) # --- Output Column --- with gr.Column(scale=1): gr.Markdown("

Results

") predicted_label = gr.Textbox(label="Predicted Material", interactive=False) confidence_score = gr.Textbox(label="Confidence", interactive=False) all_scores_label = gr.Label(num_top_classes=3, label="All Confidence Scores") # Add heatmap heatmap_gallery = gr.Gallery( label="Visualizations", columns=3, height=300 ) submit_btn = gr.Button("Classify", variant="primary") # --- Button Logic --- submit_btn.click( fn=classify_and_update, inputs=[image_input, history_state], outputs=[heatmap_gallery, predicted_label, confidence_score, all_scores_label, *history_slots, history_state] ) # ========== TAB 2: LIVE CAMERA ========== with gr.TabItem("Live IP Webcam"): json_out_live = gr.JSON(label="Prediction (top class + confidence %)") label_out_live = gr.Label(num_top_classes=3, label="Top-3 probabilities") live_feed = gr.Gallery(label="Live Feed", height=500, # Adjust to fit your page columns=1 # 1 image per row ) motion_out = gr.JSON(label="Motion Info") start_btn = gr.Button("Start Live Feed!!!") stop_btn = gr.Button("Stop Live Feed") # Start live feed (motion-triggered) start_btn.click( live_ipcam_generator, inputs=[], outputs=[json_out_live, label_out_live, live_feed, motion_out] ) # Stop button can just close the browser tab or set a global stop flag # ---- RUN THE APP ---- if __name__ == "__main__": demo.launch(inbrowser=True, share=True)