Spaces:
Sleeping
Sleeping
| 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) |