Update app.py

app.py

@@ -1,77 +1,206 @@
 import os
 import shutil
+import zipfile
+import torch
+import numpy as np
+import nibabel as nib
+import matplotlib.pyplot as plt
 import gradio as gr
-from nnunet.inference.predict import predict_from_folder
-from nnunet.paths import default_plans_identifier, network_training_output_dir, preprocessing_output_dir
+from monai.networks.nets import SwinUNETR
+from monai.inferers import sliding_window_inference
+from monai.transforms import (
+    Compose, LoadImaged, EnsureChannelFirstd, Spacingd,
+    Orientationd, ScaleIntensityd, ToTensord, MapTransform, 
+    Resized, Resize
+)
 
+# ==========================================
 # 1. CONFIGURATION
-# We have to "trick" nnU-Net into finding the model in the standard folder structure
-# because it relies on strict paths.
-
-# Set environment variables required by nnU-Net
-os.environ['nnUNet_raw_data_base'] = '/tmp/nnUNet_raw_data_base'
-os.environ['nnUNet_preprocessed'] = '/tmp/nnUNet_preprocessed'
-os.environ['RESULTS_FOLDER'] = '/tmp/nnUNet_trained_models'
-
-# Function to setup the "Fake" directory structure nnU-Net expects
-def setup_model_folder(uploaded_model_path):
-    target_folder = "/tmp/nnUNet_trained_models/nnUNet/3d_fullres/Task012_Ischemic_Stroke_TM_Fullset/nnUNetTrainerV2_DDP__nnUNetPlansv2.1/fold_0"
-    os.makedirs(target_folder, exist_ok=True)
-    
-    # Move the uploaded model to the expected location
-    # Note: nnU-Net expects 'model_best.model' and 'model_best.model.pkl'
-    # You might need to upload both or ZIP them.
-    # For this script, we assume you uploaded a ZIP of the 'fold_0' folder.
-    import zipfile
-    with zipfile.ZipFile(uploaded_model_path, 'r') as zip_ref:
-        zip_ref.extractall(target_folder)
-        
-    return target_folder
-
-# 2. INFERENCE LOGIC
-def segment_stroke(dwi, adc, model_zip):
-    # Setup Input Folder
-    input_folder = "/tmp/input_data"
-    output_folder = "/tmp/output_prediction"
-    os.makedirs(input_folder, exist_ok=True)
-    os.makedirs(output_folder, exist_ok=True)
+# ==========================================
+MODEL_FILE = "vision_agent_best.pth" 
+device = torch.device("cpu") # Use CPU for Free Tier Spaces
+
+print("--- 👁️ Initializing Vision Agent ---")
+
+# Initialize Model (Must match training architecture: 2 Channel IN -> 2 Channel OUT)
+model = SwinUNETR(
+    img_size=(96, 96, 96),
+    in_channels=2, 
+    out_channels=2, 
+    feature_size=48,
+    use_checkpoint=False, 
+    spatial_dims=3
+).to(device)
+
+# Load Weights
+if os.path.exists(MODEL_FILE):
+    try:
+        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}")
+else:
+    print(f"⚠️ Warning: {MODEL_FILE} not found. Upload it to Files.")
+
+# ==========================================
+# 2. SMART FILE HANDLER (Matches Gatekeeper Logic)
+# ==========================================
+def handle_uploads(file_paths):
+    """Flattens zips and lists into single file paths."""
+    if not file_paths: return []
+    processed_files = []
+    extract_root = "/tmp/vision_uploads"
+    if os.path.exists(extract_root): shutil.rmtree(extract_root)
+    os.makedirs(extract_root, exist_ok=True)
+
+    for path in file_paths:
+        if path.endswith(".zip"):
+            try:
+                with zipfile.ZipFile(path, 'r') as zip_ref:
+                    zip_ref.extractall(extract_root)
+                for root, _, files in os.walk(extract_root):
+                    for file in files:
+                        if not file.startswith("."): 
+                            processed_files.append(os.path.join(root, file))
+            except: pass
+        else:
+            processed_files.append(path)
+    return processed_files
+
+def find_sequences(file_list):
+    """Scans inputs to find DWI and ADC automatically."""
+    all_files = handle_uploads(file_list)
+    dwi_path = None
+    adc_path = None
     
-    # nnU-Net expects specific naming: CASE_0000.nii.gz, CASE_0001.nii.gz
-    # We rename our inputs to match what the model was trained on (Channel 0=DWI, Channel 1=ADC)
-    # Note: You must check the DeepISLES readme for channel order. Assuming DWI=0000, ADC=0001.
-    shutil.copy(dwi.name, os.path.join(input_folder, "test_0000.nii.gz"))
-    shutil.copy(adc.name, os.path.join(input_folder, "test_0001.nii.gz"))
+    print(f"Scanning {len(all_files)} files for DWI/ADC...")
     
-    # Setup Model
-    setup_model_folder(model_zip.name)
+    for path in all_files:
+        fname = os.path.basename(path).lower()
+        # Heuristic matching
+        if "dwi" in fname: dwi_path = path
+        elif "adc" in fname: adc_path = path
     
-    # Run Inference
-    # This calls the internal nnU-Net prediction function
-    # -i input, -o output, -t Task ID, -m Model Type, -f fold
-    command = f"nnUNet_predict -i {input_folder} -o {output_folder} -t 012 -m 3d_fullres -f 0"
-    os.system(command)
+    return dwi_path, adc_path
+
+# ==========================================
+# 3. PREPROCESSING PIPELINE
+# ==========================================
+class HandleMissingADCd(MapTransform):
+    def __call__(self, data):
+        # If ADC is missing (flagged in data dict), zero it out
+        if data.get("missing_adc", False):
+            data["image"][1] = data["image"][1] * 0.0 
+        return data
+
+preprocess = Compose([
+    LoadImaged(keys=["image"]),
+    EnsureChannelFirstd(keys=["image"]),
+    HandleMissingADCd(keys=["image"]),
+    ScaleIntensityd(keys=["image"]), 
+    Orientationd(keys=["image"], axcodes="RAS"),
+    Spacingd(keys=["image"], pixdim=(1.5, 1.5, 1.5), mode=("bilinear")),
+    Resized(keys=["image"], spatial_size=(96, 96, 96), mode=("bilinear")),
+    ToTensord(keys=["image"]),
+])
+
+def generate_overlay(original_dwi_path, mask_data):
+    """Visualizes the stroke slice."""
+    img = nib.load(original_dwi_path)
+    dwi_vol = img.get_fdata()
+    dwi_vol = np.rot90(dwi_vol, k=1, axes=(0, 1))
+    mask_vol = np.rot90(mask_data, k=1, axes=(0, 1))
+
+    # Find slice with largest lesion
+    slice_sums = np.sum(mask_vol, axis=(0, 1))
+    max_slice = np.argmax(slice_sums) if np.max(slice_sums) > 0 else dwi_vol.shape[2] // 2
+
+    fig, ax = plt.subplots(figsize=(5, 5))
+    ax.imshow(dwi_vol[:, :, max_slice], cmap="gray")
+    if np.max(mask_vol) > 0:
+        mask_display = np.zeros((*mask_vol.shape[:2], 4))
+        mask_display[mask_vol[:, :, max_slice] > 0] = [1, 0, 0, 0.6] # Red
+        ax.imshow(mask_display)
     
-    # Find Output
-    output_file = os.path.join(output_folder, "test.nii.gz")
-    if os.path.exists(output_file):
-        return output_file
-    else:
-        return "Error: Segmentation failed."
-
-# 3. UI
-with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 DeepISLES (SEALS Node)")
-    gr.Markdown("Upload DWI + ADC and the 'fold_0.zip' model weights.")
+    ax.axis("off")
+    output_path = "/tmp/lesion_overlay.png"
+    plt.savefig(output_path, bbox_inches='tight')
+    plt.close()
+    return output_path
+
+# ==========================================
+# 4. INFERENCE ENGINE
+# ==========================================
+def segment_stroke(file_list):
+    # 1. Find Data
+    dwi_path, adc_path = find_sequences(file_list)
     
-    with gr.Row():
-        dwi_in = gr.File(label="DWI Scan")
-        adc_in = gr.File(label="ADC Scan")
-        model_in = gr.File(label="Model Weights (Zip of fold_0 folder)")
+    if not dwi_path:
+        return {"error": "CRITICAL: No DWI sequence found."}, None
+
+    # 2. Prepare Input
+    # If ADC is missing, duplicate DWI but flag it for zero-filling
+    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)}
     
-    btn = gr.Button("Segment")
-    out = gr.File(label="Lesion Mask")
+    try:
+        # 3. Preprocess
+        tensor_data = preprocess(input_data)["image"].unsqueeze(0).to(device)
+        
+        # 4. Inference
+        with torch.no_grad():
+            output = model(tensor_data)
+            pred_mask_small = torch.argmax(output, dim=1).cpu().numpy()[0]
+
+        # 5. 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()
+        
+        # 6. Calc Volume
+        voxel_vol_mm3 = np.prod(original_img.header.get_zooms())
+        lesion_voxels = np.count_nonzero(pred_mask_original)
+        volume_ml = (lesion_voxels * voxel_vol_mm3) / 1000.0
+        
+        # 7. Report
+        risk = "LOW"
+        if volume_ml > 70: risk = "HIGH (Futile?)"
+        elif volume_ml > 50: risk = "MODERATE"
+        
+        result = {
+            "status": "SUCCESS",
+            "core_volume_ml": round(volume_ml, 2),
+            "risk_category": risk,
+            "modality_status": "DWI+ADC" if adc_path else "DWI Only (Approximation)"
+        }
+        
+        return result, generate_overlay(dwi_path, pred_mask_original)
+
+    except Exception as e:
+        return {"error": str(e)}, None
+
+# ==========================================
+# 5. GRADIO UI
+# ==========================================
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 👁️ Neuro-Triage: Vision Agent")
     
-    btn.click(segment_stroke, inputs=[dwi_in, adc_in, model_in], outputs=out)
+    with gr.Row():
+        with gr.Column():
+            file_input = gr.File(
+                label="Upload Patient Data", 
+                file_count="multiple", 
+                type="filepath"
+            )
+            btn = gr.Button("Calculate Core Volume", variant="primary")
+        
+        with gr.Column():
+            out_json = gr.JSON(label="Volumetric Report")
+            out_img = gr.Image(label="Lesion Map")
+        
+    btn.click(segment_stroke, inputs=file_input, outputs=[out_json, out_img])
 
 if __name__ == "__main__":
     demo.launch()