Update app.py

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
 import os
-import shutil
+import glob
 import zipfile
 import torch
 import numpy as np
 import nibabel as nib
 import matplotlib.pyplot as plt
-import traceback # <--- Added for debugging
-import tempfile  # <--- Added for safe temp folders
+import traceback 
+import tempfile
 from monai.networks.nets import SwinUNETR
 from monai.inferers import sliding_window_inference
 from monai.transforms import (
@@ -17,14 +17,38 @@ from monai.transforms import (
 )
 
 # ==========================================
-# 1. CONFIGURATION
+# 1. MODEL SETUP & WEIGHT EXTRACTION
 # ==========================================
-MODEL_FILE = "vision_agent_best.pth" 
 device = torch.device("cpu")
 
 print("--- 👁️ Initializing Vision Agent ---")
 
-# Initialize Model (No img_size, consistent with training)
+# --- HELPER: Find Model Weights ---
+def locate_model_weights():
+    # 1. Look for direct .pth files
+    pth_files = glob.glob("*.pth")
+    if pth_files: return pth_files[0]
+    
+    # 2. Look inside zips (e.g., if you zipped your model)
+    zips = glob.glob("*.zip")
+    for z in zips:
+        print(f"📦 Scanning zip: {z}")
+        try:
+            with zipfile.ZipFile(z, 'r') as zf:
+                # List files in zip
+                for name in zf.namelist():
+                    if name.endswith(".pth"):
+                        print(f"   -> Found weights in zip: {name}")
+                        # Extract just the weights
+                        zf.extract(name, ".")
+                        return name
+        except: pass
+        
+    return None
+
+MODEL_FILE = locate_model_weights()
+
+# Initialize Model (Swin UNETR - Matches your training)
 model = SwinUNETR(
     in_channels=2, 
     out_channels=2, 
@@ -33,17 +57,18 @@ model = SwinUNETR(
     spatial_dims=3
 ).to(device)
 
-# Load Weights with Strict Check
-if os.path.exists(MODEL_FILE):
+if MODEL_FILE and os.path.exists(MODEL_FILE):
     try:
+        print(f"🔄 Loading weights from: {MODEL_FILE}")
         state_dict = torch.load(MODEL_FILE, map_location=device)
         model.load_state_dict(state_dict)
         model.eval()
         print("✅ Model Loaded Successfully.")
     except Exception as e:
         print(f"❌ Error loading weights: {e}")
+        print("   (Are you trying to load nnU-Net/DeepISLES weights? This code expects SwinUNETR weights.)")
 else:
-    print(f"⚠️ WARNING: '{MODEL_FILE}' NOT FOUND. Upload it to the Files tab!")
+    print(f"⚠️ CRITICAL WARNING: No .pth model found. Please upload 'vision_agent_best.pth' to Files.")
 
 # ==========================================
 # 2. PREPROCESSING
@@ -59,7 +84,7 @@ preprocess = Compose([
     LoadImaged(keys=["image"]),
     EnsureChannelFirstd(keys=["image"]),
     HandleMissingADCd(keys=["image"]),
-    # Safety: If input is 4D (Time series), take first volume only
+    # Handle 4D inputs if necessary
     Lambdad(keys=["image"], func=lambda x: x if x.shape[0] == 2 else x[:, :, :, :, 0] if x.ndim==5 else x), 
     ScaleIntensityd(keys=["image"]), 
     Orientationd(keys=["image"], axcodes="RAS"),
@@ -90,7 +115,7 @@ def generate_overlay(original_dwi_path, mask_data):
         img = nib.load(original_dwi_path)
         dwi_vol = img.get_fdata()
         
-        # Orientation check (simple rot90 for display)
+        # Orientation check
         dwi_vol = np.rot90(dwi_vol, k=1, axes=(0, 1))
         mask_vol = np.rot90(mask_data, k=1, axes=(0, 1))
 
@@ -119,35 +144,35 @@ def generate_overlay(original_dwi_path, mask_data):
 # 4. INFERENCE ENGINE
 # ==========================================
 def segment_stroke(zip_file):
-    # Debug Info
     if not zip_file: return {"error": "No file received"}, None
     
     try:
-        # 1. Create Unique Temp Dir
         with tempfile.TemporaryDirectory() as extract_dir:
-            # 2. Unzip
-            with zipfile.ZipFile(zip_file, 'r') as zip_ref:
-                zip_ref.extractall(extract_dir)
-            
-            # 3. Find Data
-            dwi_path, adc_path = find_sequences(extract_dir)
+            # Handle Zip
+            if hasattr(zip_file, 'name') and zip_file.name.endswith('.zip'):
+                 with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+                    zip_ref.extractall(extract_dir)
+                 dwi_path, adc_path = find_sequences(extract_dir)
+            # Handle Raw File (if user uploads just one file by mistake)
+            else:
+                return {"error": "Please upload a .zip file containing the patient scan."}, None
             
             if not dwi_path:
                 return {"status": "FAILED", "error": "No DWI file found in zip."}, None
 
-            # 4. Prepare
+            # Prepare Input
             img_list = [dwi_path, adc_path] if adc_path else [dwi_path, dwi_path]
             input_data = {"image": img_list, "missing_adc": not bool(adc_path)}
             
-            # 5. Preprocess
+            # Preprocess
             tensor_data = preprocess(input_data)["image"].unsqueeze(0).to(device)
             
-            # 6. Inference
+            # Inference
             with torch.no_grad():
                 output = model(tensor_data)
                 pred_mask_small = torch.argmax(output, dim=1).cpu().numpy()[0]
 
-            # 7. Volume Calc (Unsquash)
+            # Volume Recalculation (Unsquash)
             original_img = nib.load(dwi_path)
             resizer = Resize(spatial_size=original_img.shape, mode="nearest")
             pred_mask_original = resizer(torch.from_numpy(pred_mask_small).unsqueeze(0))[0].numpy()
@@ -156,7 +181,7 @@ def segment_stroke(zip_file):
             voxel_vol_mm3 = np.prod(zooms)
             volume_ml = (np.count_nonzero(pred_mask_original) * voxel_vol_mm3) / 1000.0
             
-            # 8. Result
+            # Report
             risk = "LOW"
             if volume_ml > 70: risk = "HIGH (Futile?)"
             elif volume_ml > 50: risk = "MODERATE"
@@ -171,15 +196,14 @@ def segment_stroke(zip_file):
             return result, generate_overlay(dwi_path, pred_mask_original)
 
     except Exception as e:
-        # CAPTURE THE EXACT ERROR
         error_msg = traceback.format_exc()
-        print(error_msg)
-        return {"status": "CRASHED", "error": str(e), "traceback": error_msg[-200:]}, None
+        return {"status": "CRASHED", "error": str(e), "details": error_msg[-200:]}, None
 
 # ==========================================
-# 5. UI
+# 5. GRADIO UI
 # ==========================================
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+# FIX: Removed 'theme' argument
+with gr.Blocks() as demo:
     gr.Markdown("# 👁️ Neuro-Triage: Vision Agent")
     
     with gr.Row():