Update app.py

app.py

@@ -7,9 +7,11 @@ import nibabel as nib
 import traceback
 import tempfile
 import torch
+from scipy.ndimage import distance_transform_edt
+
+
+# SET ENV VARIABLES *BEFORE* IMPORTING NNUNET
 
-# 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'
@@ -19,158 +21,187 @@ os.makedirs('./tmp/nnUNet_preprocessed', exist_ok=True)
 from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
 from batchgenerators.utilities.file_and_folder_operations import join
 
-# Dynamically assign device (prevents crashes if deployed on a CPU space)
+# Model Initializations
 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
+        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',
+        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}")
+    print(f"Warning: nnU-Net failed to load. {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):
+
+# FEATURE EXTRACTION
+
+def extract_features(seg_path, high_risk_threshold_mm=5.0):
+    """Calculates metrics directly from a NIfTI mask."""
+    seg_img = nib.load(seg_path)
+    seg_data = seg_img.get_fdata()
+    spacing = seg_img.header.get_zooms()[:3]
+    voxel_volume_mm3 = spacing[0] * spacing[1] * spacing[2]
+
+    tumor_mask = (seg_data == 10)
+    artery_mask = np.isin(seg_data, [2, 8])
+    veins_mask = (seg_data == 9)
+
+    if not np.any(tumor_mask):
+        return 50.0, 0.0, 0.0, 0.0 
+
+    ttv_mm3 = np.sum(tumor_mask) * voxel_volume_mm3
+
+    dist_artery = distance_transform_edt(~artery_mask, sampling=spacing) if np.any(artery_mask) else np.full(seg_data.shape, 50.0)
+    dist_veins = distance_transform_edt(~veins_mask, sampling=spacing) if np.any(veins_mask) else np.full(seg_data.shape, 50.0)
+
+    tumor_artery_distances = dist_artery[tumor_mask]
+    tumor_vein_distances = dist_veins[tumor_mask]
+
+    mac = np.min(tumor_artery_distances)
+    aev_mm3 = np.sum(tumor_artery_distances <= high_risk_threshold_mm) * voxel_volume_mm3
+    vev_mm3 = np.sum(tumor_vein_distances <= high_risk_threshold_mm) * voxel_volume_mm3
+
+    return mac, vev_mm3, aev_mm3, ttv_mm3
+
+def calculate_risk_and_format(mask_path, no_chemo, pd_surgery):
+    """Runs extraction, ML prediction, and formats UI output."""
+    mac, vev, aev, ttv = extract_features(mask_path)
+
+    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 # Fallback dummy
+
+    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)
+
+
+# PIPELINE HANDLERS
+
+def process_raw_ct(nifti_file, no_chemo, pd_surgery):
+    """Panel 1: Runs heavy nnU-Net segmentation first."""
     if nifti_file is None:
-        return gr.update(value="**Error:** Please upload a NIfTI file.", visible=True), *[gr.update(value="")]*4, gr.update(value="")
+        return gr.update(value="**Error:** Please upload a Raw CT 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
+            # Ensure _0000 suffix
             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.")
+                raise RuntimeError("nnU-Net is not initialized.")
 
-            # Predict from the temporary input folder to the temporary output folder
+            # Run Segmentation
             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
+                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   
+                raise FileNotFoundError("nnU-Net failed to generate mask.")
 
-            # 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 calculate_risk_and_format(output_mask_path, no_chemo, pd_surgery)
 
-            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"**nnU-Net 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="")
 
+def process_presegmented_mask(mask_file, no_chemo, pd_surgery):
+    """Panel 2: Skips nnU-Net, processes mask directly."""
+    if mask_file is None:
+        return gr.update(value="**Error:** Please upload a pre-segmented Mask NIfTI file.", visible=True), *[gr.update(value="")]*4, gr.update(value="")
+    
+    try:
+        # Directly pass the uploaded mask to the extraction logic
+        return calculate_risk_and_format(mask_file.name, no_chemo, pd_surgery)
     except Exception as e:
-        error_msg = f"**Pipeline Error:** {str(e)}\n\n```\n{traceback.format_exc()}\n```"
+        error_msg = f"**Fast 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
 
+# GRADIO UI LAYOUT
 
 with gr.Blocks() as demo:
     gr.Markdown("# Preoperative R1 Resection Risk Predictor")
     
     with gr.Row():
+        # LEFT COLUMN: Inputs (Tabs + Clinical Data)
         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")
+            with gr.Tabs():
+                # TAB 1: Slow / Full Pipeline
+                with gr.TabItem("Raw CT (Runs nnU-Net)"):
+                    gr.Markdown("*Note: Segmentation may take >1 hour on CPU.*")
+                    file_input_ct = gr.File(label="Upload Raw CT Scan (.nii.gz)", file_count="single")
+                    analyze_ct_btn = gr.Button("Segment & Analyze", variant="primary")
+                
+                # TAB 2: Fast Pipeline
+                with gr.TabItem("Pre-segmented Mask (Fast)"):
+                    gr.Markdown("*Bypass AI segmentation. Upload your own mask.*")
+                    file_input_mask = gr.File(label="Upload Segmentation Mask (.nii.gz)", file_count="single")
+                    analyze_mask_btn = gr.Button("Analyze Mask", variant="primary")
+            
+            gr.Markdown("---")
+            gr.Markdown("### Clinical Metadata (Required)")
             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")
-            
+        # RIGHT COLUMN: Shared Outputs
         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)
+                mac_out = gr.Textbox(label="MAC (mm)", interactive=False)
+                vev_out = gr.Textbox(label="VEV (mm³)", interactive=False)
+                aev_out = gr.Textbox(label="AEV (mm³)", interactive=False)
+                ttv_out = gr.Textbox(label="TTV (mm³)", interactive=False)
             
             gr.Markdown("### Risk Assessment")
-            risk_out = gr.HTML(value="<i>Run analysis to see predicted risk...</i>")
+            risk_out = gr.HTML(value="<i>Upload a file and run analysis...</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]
+    # Connect buttons to their respective functions, but update the SAME output boxes
+    outputs_list = [error_box, mac_out, vev_out, aev_out, ttv_out, risk_out]
+    
+    analyze_ct_btn.click(
+        fn=process_raw_ct,
+        inputs=[file_input_ct, no_chemo, pd_surgery],
+        outputs=outputs_list
+    )
+    
+    analyze_mask_btn.click(
+        fn=process_presegmented_mask,
+        inputs=[file_input_mask, no_chemo, pd_surgery],
+        outputs=outputs_list
     )
 
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0", 
-        server_port=7860,
-        theme=gr.themes.Base() 
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)