Blame for app.py

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 scipy.ndimage import distance_transform_edt


# SET ENV VARIABLES *BEFORE* IMPORTING NNUNET

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)

from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
from batchgenerators.utilities.file_and_folder_operations import join

# Model Initializations
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
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
    )
    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',
    )
except Exception as e:
    predictor = None
    print(f"Warning: nnU-Net failed to load. {e}")

try:
    risk_model = joblib.load("model.joblib") 
except FileNotFoundError:
    risk_model = None



# 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."""
    # 1. Extract raw metrics
    mac, vev, aev, ttv = extract_features(mask_path)

    # Convert textboxes safely
    try:
        chemo_val = int(no_chemo)
        pd_val = int(pd_surgery)
    except ValueError:
        error_msg = "**Error:** Clinical metadata fields must be numbers (1 or 0)."
        return (gr.update(value=error_msg, visible=True), "", "", "", "", "")

    # 2. Hardcode the Means and Stds from your training data
    # Extracted from your means/stds arrays based on their index positions
    TTV_MEAN, TTV_STD = 8980.23827, 8281.19811
    MAC_MEAN, MAC_STD = 8.53567157, 8.92353600
    AEV_MEAN, AEV_STD = 155.784393, 367.152035
    VEV_MEAN, VEV_STD = 576.025386, 559.944824

    # 3. Apply Z-score Normalization
    scaled_ttv = (ttv - TTV_MEAN) / TTV_STD
    scaled_mac = (mac - MAC_MEAN) / MAC_STD
    scaled_aev = (aev - AEV_MEAN) / AEV_STD
    scaled_vev = (vev - VEV_MEAN) / VEV_STD

    # 4. Predict Risk
    if risk_model is not None:
        # STRICT ORDER MATCHING your column_order:
        # ['Chemo_0', 'Surgery_1', 'Total_Tumor_Volume_mm3_arterial', 'Min_Distance_mm_arterial', 'High_Risk_Volume_mm3_5mm_arterial', 'High_Risk_Volume_mm3_5mm_venous']
        features = np.array([[
            chemo_val,   # Chemo_0
            pd_val,      # Surgery_1
            scaled_ttv,  # Total_Tumor_Volume_mm3_arterial
            scaled_mac,  # Min_Distance_mm_arterial
            scaled_aev,  # High_Risk_Volume_mm3_5mm_arterial
            scaled_vev   # High_Risk_Volume_mm3_5mm_venous
        ]])
        
        prob_r1 = risk_model.predict_proba(features)[0][1] 
    else:
        prob_r1 = 0.72 # Fallback dummy

    # 5. Format Outputs
    risk_percentage = prob_r1
    risk_color = "#C55A11" if risk_percentage > 0.5 else "#B89036"
    risk_display = f"<h2 style='color: {risk_color}; margin-top: 0;'>Predicted probability of R+ resection: {risk_percentage:.2f}</h2>"

    # Notice we still display the RAW un-normalized values to the user in the UI, 
    # but the model calculates risk using the SCALED values.
    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 Raw CT NIfTI file.", visible=True), *[gr.update(value="")]*4, gr.update(value="")
    
    try:
        with tempfile.TemporaryDirectory() as temp_in_dir, tempfile.TemporaryDirectory() as temp_out_dir:
            original_path = nifti_file.name
            base_name = os.path.basename(original_path)

            # 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")
            else:
                new_base = base_name

            input_file_path = os.path.join(temp_in_dir, new_base)
            shutil.copy(original_path, input_file_path)

            if predictor is None:
                raise RuntimeError("nnU-Net is not initialized.")

            # 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
            )

            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 mask.")

            return calculate_risk_and_format(output_mask_path, no_chemo, pd_surgery)

    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"**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

with gr.Blocks() as demo:
    gr.Markdown("# Preoperative PDAC R+ Resection Risk Predictor")
    
    with gr.Row():
        # LEFT COLUMN: Inputs (Tabs + Clinical Data)
        with gr.Column(scale=1):
            
            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)")
            
        # 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="Minimum Arterial Clearance (mm)", interactive=False)
                vev_out = gr.Textbox(label="Venous Encasement Volume (mm³)", interactive=False)
                aev_out = gr.Textbox(label="Arterial Encasement Volume (mm³)", interactive=False)
                ttv_out = gr.Textbox(label="Total Tumor Volume (mm³)", interactive=False)
            
            gr.Markdown("### Risk Assessment")
            risk_out = gr.HTML(value="<i>Upload a file and run analysis...</i>")

    # 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, share=True)