Update app.py

app.py

@@ -14,11 +14,9 @@ print(f"--- System Boot: Using {device} ---")
 
 # --- LOAD VISUAL SYSTEM (YOLO) ---
 try:
-    # Try loading your custom trained model first
     yolo_model = YOLO("best.pt")
     print("✅ Visual System: Custom EAGLE A7 Model Loaded")
 except:
-    print("⚠️ Visual System: Custom model not found. Loading standard YOLOv11n...")
     yolo_model = YOLO("yolo11n.pt")
 
 # --- LOAD THERMAL SYSTEM (ResNet-18) ---
@@ -29,139 +27,130 @@ def get_thermal_model():
     return model
 
 thermal_model = get_thermal_model().to(device)
-
 MODEL_PATH = "thermal_landmine_scanner.pth"
+
 try:
     thermal_model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
     thermal_model.eval()
     print(f"✅ Thermal System: Loaded {MODEL_PATH}")
 except Exception as e:
     print(f"❌ CRITICAL ERROR: Could not load thermal model. {e}")
-    print("   Make sure 'thermal_landmine_scanner.pth' is in this folder!")
 
 # ==========================================
-# 2. CORE PROCESSING FUNCTIONS
+# 2. SETUP GRAD-CAM (THE "X-RAY" HOOK)
+# ==========================================
+# We need to steal the features from inside the model while it thinks
+features_blob = []
+def hook_feature(module, input, output):
+    features_blob.clear() # Clear old data
+    features_blob.append(output.data.cpu().numpy())
+
+# Attach the spy hook to the last layer (Layer 4)
+thermal_model.layer4.register_forward_hook(hook_feature)
+
+# Get weights from the final decision layer
+params = list(thermal_model.parameters())
+weight_softmax = params[-2].data.cpu().numpy() # The weights connecting features to "Mine/Safe"
+
+# ==========================================
+# 3. PROCESSING FUNCTIONS
 # ==========================================
 
 def run_visual_detection(image):
-    """
-    Daytime RGB Detection using YOLOv11
-    """
-    if image is None: return None, "Waiting for video feed..."
-    
-    # YOLO handles preprocessing internally
-    results = yolo_model.predict(image, conf=0.40) # 40% Confidence Threshold
-    
-    # Render detections
-    annotated_frame = results[0].plot()
-    
-    # Status Logic
-    count = len(results[0].boxes)
-    if count > 0:
-        status = f"⚠️ DANGER: {count} Threat(s) Identified"
-    else:
-        status = "✅ Sector Clear"
-        
-    return annotated_frame, status
+    if image is None: return None, "Waiting for feed..."
+    results = yolo_model.predict(image, conf=0.40)
+    return results[0].plot(), f"Objects Detected: {len(results[0].boxes)}"
 
 def run_thermal_scan(image):
-    """
-    Nighttime Thermal Analysis using ResNet-18
-    Strict preprocessing to match training data.
-    """
-    if image is None: return "N/A", "No Signal"
-
-    # --- STEP 1: FORCE GRAYSCALE ---
-    # Gradio sends RGB. We need Grayscale.
-    # We use RGB2GRAY because Gradio is RGB by default.
+    if image is None: return None, "No Signal", "N/A"
+
+    # --- PREPROCESSING (Standard) ---
     if len(image.shape) == 3:
         gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
     else:
         gray = image
 
-    # --- STEP 2: CLAHE ENHANCEMENT ---
-    # This MUST match the training/debug script exactly.
     clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
     enhanced_img = clahe.apply(gray)
-    
-    # --- STEP 3: RESIZE ---
     resized = cv2.resize(enhanced_img, (224, 224))
-    
-    # --- STEP 4: NORMALIZE TO 0.0 - 1.0 ---
-    # Convert to float32 and divide by 255
     normalized_img = resized.astype(np.float32) / 255.0
     
-    # --- STEP 5: TENSOR CONVERSION ---
-    # Shape: [1, 1, 224, 224]
     tensor = torch.from_numpy(normalized_img).float().unsqueeze(0).unsqueeze(0)
     tensor = tensor.to(device)
     
-    # --- STEP 6: INFERENCE ---
+    # --- INFERENCE ---
     with torch.no_grad():
         output = thermal_model(tensor)
         prob = torch.sigmoid(output).item()
         
-    # --- DEBUG LOGGING (Check your terminal!) ---
-    # This proves if the "App" sees what the "Debug Script" saw.
-    print(f"🔎 Thermal Scan | Input Mean: {normalized_img.mean():.4f} | Prob: {prob:.4f}")
-
-    # --- STEP 7: DECISION LOGIC ---
-    # Based on your perfect debug run (Mine=0.99, Safe=0.00), 
-    # we use a standard 0.50 threshold.
-    THRESHOLD = 0.50
+    # --- GENERATE HEATMAP (Explainable AI) ---
+    # 1. Get the features captured by our hook [1, 512, 7, 7]
+    feature_data = features_blob[0]
     
-    is_mine = prob > THRESHOLD
+    # 2. Calculate the "Attention Map"
+    cam = np.zeros((7, 7), dtype=np.float32)
+    # Use the weights for the "Mine" class to weight the features
+    for i, w in enumerate(weight_softmax[0]): 
+        cam += w * feature_data[0, i, :, :]
+        
+    # 3. Process the CAM
+    cam = np.maximum(cam, 0) # ReLU (Remove negative influence)
+    cam = cv2.resize(cam, (224, 224)) # Resize to image size
+    cam = cam - np.min(cam)
+    if np.max(cam) != 0:
+        cam = cam / np.max(cam) # Normalize 0-1
+        
+    # 4. Colorize
+    heatmap = cv2.applyColorMap(np.uint8(255*cam), cv2.COLORMAP_JET)
     
+    # 5. Overlay on original (Grayscale -> RGB)
+    orig_rgb = cv2.cvtColor(resized, cv2.COLOR_GRAY2RGB)
+    
+    # Mix: 60% Original Image + 40% Heatmap
+    # If prob is low (Safe), we show less heatmap so it doesn't look scary
+    intensity = 0.5 if prob > 0.5 else 0.2 
+    overlay = cv2.addWeighted(orig_rgb, 1-intensity, heatmap, intensity, 0)
+    
+    # --- DECISION ---
+    is_mine = prob > 0.50
     if is_mine:
-        label = "🔴 MINE DETECTED"
-        color_hex = "red"
-        conf_text = f"THREAT CONFIDENCE: {prob*100:.1f}%"
+        label = f"<h2 style='color: red; text-align: center;'>🔴 MINE DETECTED</h2>"
+        conf_text = f"CONFIDENCE: {prob*100:.1f}%"
     else:
-        label = "🟢 SAFE SOIL"
-        color_hex = "green"
-        conf_text = f"Safety Confidence: {(1-prob)*100:.1f}%"
+        label = f"<h2 style='color: green; text-align: center;'>🟢 SAFE SOIL</h2>"
+        conf_text = f"Risk Level: {prob*100:.1f}%"
 
-    # Return HTML for colored text
-    html_output = f"<h2 style='color: {color_hex}; text-align: center;'>{label}</h2>"
-    return html_output, conf_text
+    return overlay, label, conf_text
 
 # ==========================================
-# 3. DASHBOARD UI (Gradio)
+# 4. DASHBOARD UI
 # ==========================================
-custom_css = """
-.gradio-container {background-color: #1e1e1e; color: white}
-"""
+custom_css = ".gradio-container {background-color: #1e1e1e; color: white}"
 
 with gr.Blocks(css=custom_css, title="EAGLE A7 Mission Control") as demo:
     gr.Markdown("# 🦅 EAGLE A7: Autonomous Demining Interface")
-    gr.Markdown("**Status:** System Online | **AI Engines:** YOLOv11 + Thermal ResNet-18")
     
     with gr.Tabs():
-        # --- TAB 1: VISUAL ---
         with gr.TabItem("☀️ Daytime Vision"):
             with gr.Row():
-                with gr.Column():
-                    vis_input = gr.Image(label="Drone Camera Feed", type="numpy")
-                    vis_btn = gr.Button("SCAN SECTOR", variant="primary")
-                with gr.Column():
-                    vis_output = gr.Image(label="AI Analysis")
-                    vis_status = gr.Textbox(label="Mission Status", interactive=False)
-            
+                vis_input = gr.Image(label="Input", type="numpy")
+                vis_output = gr.Image(label="YOLO Detections")
+            vis_btn = gr.Button("SCAN", variant="primary")
+            vis_status = gr.Textbox(label="Status")
             vis_btn.click(run_visual_detection, inputs=vis_input, outputs=[vis_output, vis_status])
 
-        # --- TAB 2: THERMAL ---
-        with gr.TabItem("🌙 Night Vision (Thermal)"):
+        with gr.TabItem("🌙 Night Vision (X-Ray)"):
+            gr.Markdown("### Thermal Anomaly Localization")
             with gr.Row():
                 with gr.Column():
-                    therm_input = gr.Image(label="Thermal Sensor Feed (Grayscale)", type="numpy")
-                    therm_btn = gr.Button("ANALYZE HEAT SIGNATURE", variant="stop")
+                    therm_input = gr.Image(label="Thermal Feed", type="numpy")
+                    therm_btn = gr.Button("ANALYZE & LOCATE", variant="stop")
                 with gr.Column():
-                    # We use HTML for big colored alerts
+                    therm_output = gr.Image(label="Target Localization (Heatmap)")
                     therm_label = gr.HTML(label="Result")
-                    therm_conf = gr.Textbox(label="Sensor Telemetry", interactive=False)
+                    therm_conf = gr.Textbox(label="Telemetry")
             
-            therm_btn.click(run_thermal_scan, inputs=therm_input, outputs=[therm_label, therm_conf])
+            therm_btn.click(run_thermal_scan, inputs=therm_input, outputs=[therm_output, therm_label, therm_conf])
 
-# Launch
-print("--- Launching Dashboard ---")
+print("--- Launching Dashboard with X-Ray Vision ---")
 demo.launch(server_name="0.0.0.0", share=True)