Update app.py

app.py

@@ -11,21 +11,17 @@ import subprocess
 # ==========================================
 # 1. DEEPISLES SETUP (Environment Trickery)
 # ==========================================
-# nnU-Net requires specific env vars to find models.
-# We set them to temporary folders.
 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'
 
-# Define the rigid structure nnU-Net expects for the model
 TASK_NAME = "Task012_Ischemic_Stroke_TM_Fullset"
 TRAINER = "nnUNetTrainerV2_DDP__nnUNetPlansv2.1"
 MODEL_PATH_INTERNAL = os.path.join(os.environ['RESULTS_FOLDER'], "nnUNet", "3d_fullres", TASK_NAME, TRAINER, "fold_0")
 
 print("--- 🧠 Initializing DeepISLES Engine ---")
 
-# 1. Install Weights (Run once on startup)
-UPLOADED_WEIGHTS = "fold_0.zip" # <--- Expects this file in Space
+UPLOADED_WEIGHTS = "fold_0.zip" 
 
 if os.path.exists(UPLOADED_WEIGHTS):
     if not os.path.exists(MODEL_PATH_INTERNAL):
@@ -44,7 +40,6 @@ else:
 # 2. FILE HANDLING
 # ==========================================
 def find_sequences(extract_dir):
-    """Finds DWI (b1000) and ADC files."""
     dwi = None
     adc = None
     for root, _, files in os.walk(extract_dir):
@@ -55,29 +50,26 @@ def find_sequences(extract_dir):
     return dwi, adc
 
 def generate_overlay(dwi_path, mask_path):
-    """Creates visual check."""
     try:
         img = nib.load(dwi_path)
         mask = nib.load(mask_path)
         dwi = img.get_fdata()
         seg = mask.get_fdata()
-
-        # Rotate for display
+        
         dwi = np.rot90(dwi, k=1, axes=(0, 1))
         seg = np.rot90(seg, k=1, axes=(0, 1))
-
-        # Find best slice
+        
         sums = np.sum(seg, axis=(0, 1))
         idx = np.argmax(sums) if np.max(sums) > 0 else dwi.shape[2]//2
-
+        
         fig, ax = plt.subplots(figsize=(5, 5))
         ax.imshow(dwi[:, :, idx], cmap="gray")
         if np.max(seg) > 0:
             overlay = np.zeros((*seg.shape[:2], 4))
-            overlay[seg[:, :, idx] > 0] = [1, 0, 0, 0.6] # Red
+            overlay[seg[:, :, idx] > 0] = [1, 0, 0, 0.6] 
             ax.imshow(overlay)
         ax.axis("off")
-
+        
         out_path = "/tmp/overlay.png"
         plt.savefig(out_path, bbox_inches='tight')
         plt.close()
@@ -90,34 +82,28 @@ def generate_overlay(dwi_path, mask_path):
 # ==========================================
 def segment_stroke(zip_file):
     if not zip_file: return {"error": "No file uploaded"}, None
-
-    # Create temp workspace
+    
     work_dir = tempfile.mkdtemp()
     input_dir = os.path.join(work_dir, "input")
     output_dir = os.path.join(work_dir, "output")
     os.makedirs(input_dir, exist_ok=True)
     os.makedirs(output_dir, exist_ok=True)
-
+    
     try:
-        # 1. Extract Patient
-        # Gradio 5 passes file path as string, Gradio 4 passes object
-        # We handle both safely
+        # Handle Gradio 5 file path string
         file_path = zip_file.name if hasattr(zip_file, 'name') else zip_file
         
         with zipfile.ZipFile(file_path, 'r') as z:
             z.extractall(work_dir)
-
-        # 2. Find Sequences
+            
         dwi, adc = find_sequences(work_dir)
         if not dwi or not adc:
             return {"status": "FAILED", "error": "DWI or ADC missing in zip"}, None
-
-        # 3. Format for nnU-Net (Crucial Step!)
-        # nnU-Net expects: CASE_0000.nii.gz (DWI) and CASE_0001.nii.gz (ADC)
+            
+        # Format for nnU-Net: 0000=DWI, 0001=ADC
         shutil.copy(dwi, os.path.join(input_dir, "patient_0000.nii.gz"))
         shutil.copy(adc, os.path.join(input_dir, "patient_0001.nii.gz"))
-
-        # 4. Run Prediction Command
+        
         print("🚀 Running DeepISLES Inference...")
         cmd = [
             "nnUNet_predict",
@@ -128,35 +114,31 @@ def segment_stroke(zip_file):
             "-f", "0",
             "--save_npz" 
         ]
-
-        # We use subprocess to run the CLI command
+        
         process = subprocess.run(cmd, capture_output=True, text=True)
-
+        
         if process.returncode != 0:
             print("NNUNET ERROR:", process.stderr)
             return {"status": "FAILED", "error": "Model inference failed. See logs."}, None
-
-        # 5. Process Result
+            
         result_file = os.path.join(output_dir, "patient.nii.gz")
         if os.path.exists(result_file):
-            # Calculate Volume
             mask_img = nib.load(result_file)
             mask_data = mask_img.get_fdata()
             zooms = mask_img.header.get_zooms()
             vol_ml = (np.sum(mask_data > 0) * np.prod(zooms)) / 1000.0
-
+            
             risk = "LOW"
             if vol_ml > 70: risk = "HIGH (Futile Recanalization?)"
             elif vol_ml > 50: risk = "MODERATE"
-
+            
             report = {
                 "status": "SUCCESS",
                 "core_volume_ml": round(vol_ml, 2),
                 "risk_assessment": risk,
                 "model_used": "DeepISLES (nnU-Net Ensembled)"
             }
-
-            # Visualize
+            
             overlay = generate_overlay(dwi, result_file)
             return report, overlay
         else:
@@ -168,13 +150,12 @@ def segment_stroke(zip_file):
         shutil.rmtree(work_dir, ignore_errors=True)
 
 # ==========================================
-# 4. UI (THEME REMOVED)
+# 4. UI
 # ==========================================
-# FIX: Removed theme argument here to prevent crash
 with gr.Blocks() as demo:
     gr.Markdown("# 👁️ Neuro-Triage: Vision Agent")
     gr.Markdown("Upload Patient ZIP (DWI + ADC). Calculates Infarct Core Volume.")
-
+    
     with gr.Row():
         with gr.Column():
             inp = gr.File(label="Patient Data (.zip)", type="filepath")
@@ -182,7 +163,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             out_json = gr.JSON(label="Clinical Report")
             out_img = gr.Image(label="Lesion Map")
-
+            
     btn.click(segment_stroke, inputs=inp, outputs=[out_json, out_img])
 
 if __name__ == "__main__":