Spaces:
Sleeping
Sleeping
| import os | |
| import shutil | |
| import zipfile | |
| import gradio as gr | |
| import nibabel as nib | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import tempfile | |
| import subprocess | |
| # ========================================== | |
| # 1. DEEPISLES SETUP (Environment Trickery) | |
| # ========================================== | |
| 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' | |
| 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 ---") | |
| UPLOADED_WEIGHTS = "fold_0.zip" | |
| if os.path.exists(UPLOADED_WEIGHTS): | |
| if not os.path.exists(MODEL_PATH_INTERNAL): | |
| print(f"📦 Installing {UPLOADED_WEIGHTS}...") | |
| os.makedirs(MODEL_PATH_INTERNAL, exist_ok=True) | |
| try: | |
| with zipfile.ZipFile(UPLOADED_WEIGHTS, 'r') as zip_ref: | |
| zip_ref.extractall(MODEL_PATH_INTERNAL) | |
| print("✅ DeepISLES Weights Installed.") | |
| except Exception as e: | |
| print(f"❌ Error extracting weights: {e}") | |
| else: | |
| print(f"⚠️ CRITICAL: '{UPLOADED_WEIGHTS}' not found. Please upload it to the Files tab.") | |
| # ========================================== | |
| # 2. FILE HANDLING | |
| # ========================================== | |
| def find_sequences(extract_dir): | |
| dwi = None | |
| adc = None | |
| for root, _, files in os.walk(extract_dir): | |
| for f in files: | |
| name = f.lower() | |
| if "dwi" in name and ".nii" in name: dwi = os.path.join(root, f) | |
| elif "adc" in name and ".nii" in name: adc = os.path.join(root, f) | |
| return dwi, adc | |
| def generate_overlay(dwi_path, mask_path): | |
| try: | |
| img = nib.load(dwi_path) | |
| mask = nib.load(mask_path) | |
| dwi = img.get_fdata() | |
| seg = mask.get_fdata() | |
| dwi = np.rot90(dwi, k=1, axes=(0, 1)) | |
| seg = np.rot90(seg, k=1, axes=(0, 1)) | |
| 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] | |
| ax.imshow(overlay) | |
| ax.axis("off") | |
| out_path = "/tmp/overlay.png" | |
| plt.savefig(out_path, bbox_inches='tight') | |
| plt.close() | |
| return out_path | |
| except: | |
| return None | |
| # ========================================== | |
| # 3. INFERENCE ENGINE (nnU-Net Wrapper) | |
| # ========================================== | |
| def segment_stroke(zip_file): | |
| if not zip_file: return {"error": "No file uploaded"}, None | |
| 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: | |
| # 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) | |
| dwi, adc = find_sequences(work_dir) | |
| if not dwi or not adc: | |
| return {"status": "FAILED", "error": "DWI or ADC missing in zip"}, None | |
| # 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")) | |
| print("🚀 Running DeepISLES Inference...") | |
| cmd = [ | |
| "nnUNet_predict", | |
| "-i", input_dir, | |
| "-o", output_dir, | |
| "-t", "012", | |
| "-m", "3d_fullres", | |
| "-f", "0", | |
| "--save_npz" | |
| ] | |
| 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 | |
| result_file = os.path.join(output_dir, "patient.nii.gz") | |
| if os.path.exists(result_file): | |
| 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)" | |
| } | |
| overlay = generate_overlay(dwi, result_file) | |
| return report, overlay | |
| else: | |
| return {"status": "FAILED", "error": "No output mask generated"}, None | |
| except Exception as e: | |
| return {"status": "CRASHED", "error": str(e)}, None | |
| finally: | |
| shutil.rmtree(work_dir, ignore_errors=True) | |
| # ========================================== | |
| # 4. UI | |
| # ========================================== | |
| 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") | |
| btn = gr.Button("Calculate Core Volume", variant="primary") | |
| 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__": | |
| demo.launch() |