Create app.py

app.py

@@ -0,0 +1,171 @@
+import gradio as gr
+import os
+import shutil
+import joblib
+import numpy as np
+import nibabel as nib
+import traceback
+import tempfile
+import torch
+from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
+from batchgenerators.utilities.file_and_folder_operations import join
+
+# GLOBAL SETUP & MODEL INITIALIZATION
+# Set dummy environment variables to prevent nnU-Net path errors in HF Spaces
+os.environ['nnUNet_raw'] = './tmp/nnUNet_raw'
+os.environ['nnUNet_results'] = './model_weights' 
+os.environ['nnUNet_preprocessed'] = './tmp/nnUNet_preprocessed'
+os.makedirs('./tmp/nnUNet_raw', exist_ok=True)
+os.makedirs('./tmp/nnUNet_preprocessed', exist_ok=True)
+
+# Dynamically assign device (prevents crashes if deployed on a CPU space)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Initializing nnU-Net on device: {device}")
+
+try:
+    predictor = nnUNetPredictor(
+        tile_step_size=0.5,
+        use_gaussian=True,
+        use_mirroring=True,
+        perform_everything_on_device=True,
+        device=device,
+        verbose=False,
+        verbose_preprocessing=False,
+        allow_tqdm=True
+    )
+
+    # UPDATE THIS PATH: Ensure this matches the folder structure in your HF Space repository
+    model_folder = './model_weights/Dataset998_PanTSTen/nnUNetTrainer__nnUNetPlans__3d_fullres'
+    
+    predictor.initialize_from_trained_model_folder(
+        model_folder,
+        use_folds=('all',),
+        checkpoint_name='checkpoint_best.pth',
+    )
+    print("nnU-Net loaded successfully.")
+except Exception as e:
+    predictor = None
+    print(f"Warning: nnU-Net failed to load. Ensure 'model_weights' folder exists. Error: {e}")
+
+try:
+    risk_model = joblib.load("model.joblib") 
+except FileNotFoundError:
+    risk_model = None
+
+# MAIN PROCESSING PIPELINE
+
+def process_scan(nifti_file, no_chemo, pd_surgery):
+    if nifti_file is None:
+        return gr.update(value="**Error:** Please upload a NIfTI file.", visible=True), *[gr.update(value="")]*4, gr.update(value="")
+    
+    try:
+        # Use Python's tempfile to handle temporary folders automatically
+        with tempfile.TemporaryDirectory() as temp_in_dir, tempfile.TemporaryDirectory() as temp_out_dir:
+            
+            # File Validation & Renaming
+            original_path = nifti_file.name
+            base_name = os.path.basename(original_path)
+
+            # Enforce nnU-Net's _0000 naming convention
+            if not base_name.endswith("_0000.nii.gz") and base_name.endswith(".nii.gz"):
+                new_base = base_name.replace(".nii.gz", "_0000.nii.gz")
+            elif not base_name.endswith("_0000.nii") and base_name.endswith(".nii"):
+                new_base = base_name.replace(".nii", "_0000.nii")
+            else:
+                new_base = base_name
+
+            # Copy the single uploaded file into the temporary input directory
+            input_file_path = os.path.join(temp_in_dir, new_base)
+            shutil.copy(original_path, input_file_path)
+
+            # Segmentation (nnU-Net)
+            if predictor is None:
+                raise RuntimeError("nnU-Net predictor was not initialized properly at startup.")
+
+            # Predict from the temporary input folder to the temporary output folder
+            predictor.predict_from_files(
+                temp_in_dir,
+                temp_out_dir,
+                save_probabilities=False, 
+                overwrite=False,
+                num_processes_preprocessing=2, 
+                num_processes_segmentation_export=2, 
+                folder_with_segs_from_prev_stage=None,
+                num_parts=1, 
+                part_id=0
+            )
+
+            output_base = new_base.replace("_0000", "")
+            output_mask_path = os.path.join(temp_out_dir, output_base)
+
+            if not os.path.exists(output_mask_path):
+                raise FileNotFoundError("nnU-Net failed to generate the output mask.")
+
+            # --- Step 3: Feature Extraction ---
+            # TODO: Add your custom extraction logic here
+            # mask_img = nib.load(output_mask_path)
+            # mask_data = mask_img.get_fdata()
+            # mac, vev, aev, ttv = extract_features(mask_data, original_scan_data)
+            
+            # Dummy values for UI demonstration
+            mac, vev, aev, ttv = 18.4, 62.0, 4.2, 1.8   
+
+            # Risk Prediction
+            if risk_model is not None:
+                features = np.array([[mac, vev, aev, ttv, int(no_chemo), int(pd_surgery)]])
+                prob_r1 = risk_model.predict_proba(features)[0][1] 
+            else:
+                prob_r1 = 0.72 # Dummy value
+
+            # Format Output
+            risk_percentage = prob_r1 * 100
+            risk_color = "red" if risk_percentage > 50 else "green"
+            risk_display = f"<h2 style='color: {risk_color}; margin-top: 0;'>Predicted probability of R1 resection: {risk_percentage:.1f}%</h2>"
+
+            return (
+                gr.update(value="", visible=False), 
+                f"{mac:.1f}", f"{vev:.1f}", f"{aev:.1f}", f"{ttv:.1f}", 
+                risk_display
+            )
+
+    except Exception as e:
+        error_msg = f"**Pipeline Error:** {str(e)}\n\n```\n{traceback.format_exc()}\n```"
+        return gr.update(value=error_msg, visible=True), *[gr.update(value="")]*4, gr.update(value="")
+
+# GRADIO UI LAYOUT
+
+
+with gr.Blocks(theme=gr.themes.Base()) as demo:
+    gr.Markdown("# Preoperative R1 Resection Risk Predictor")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            file_input = gr.File(label="Drop your NIfTI file here", file_types=[".nii.gz", ".nii"], file_count="single")
+            
+            gr.Markdown("### Clinical Metadata")
+            no_chemo = gr.Checkbox(label="Absence of Neoadjuvant Chemo (1/0)")
+            pd_surgery = gr.Checkbox(label="Patient underwent PD (Pancreaticoduodenectomy) (1/0)")
+            
+            analyze_btn = gr.Button("Run Analysis", variant="primary")
+            
+        with gr.Column(scale=2):
+            error_box = gr.Markdown(visible=False)
+            
+            gr.Markdown("### Segmentation Metrics")
+            with gr.Row():
+                mac_out = gr.Textbox(label="MAC", interactive=False)
+                vev_out = gr.Textbox(label="VEV", interactive=False)
+                aev_out = gr.Textbox(label="AEV", interactive=False)
+                ttv_out = gr.Textbox(label="TTV", interactive=False)
+            
+            gr.Markdown("### Risk Assessment")
+            risk_out = gr.HTML(value="<i>Run analysis to see predicted risk...</i>")
+
+    analyze_btn.click(
+        fn=process_scan,
+        inputs=[file_input, no_chemo, pd_surgery],
+        outputs=[error_box, mac_out, vev_out, aev_out, ttv_out, risk_out]
+    )
+
+if __name__ == "__main__":
+    demo.launch()