From 08e0612d6029ad871d5f6bf25fcbecf944b78ebf Mon Sep 17 00:00:00 2001
From: Aditya Sanjeev <adityasanjeev28@gmail.com>
Date: Tue, 11 Aug 2026 23:15:47 -0700
Subject: [PATCH] ePAI: GPU-resident export path (lossless ~1.85x) via
 EPAI_SCRIPT_PATH hook

Profiling showed ePAI's export/convert is ~158s of CPU resampling of the 26-channel
logits back to native res -- as costly as inference. This adds scripts/epai_predict.py
(+ .sh wrapper) that reuses the fork's own export_prediction_from_logits (seg AND the
PDAC/cyst/PNET findings CSV) but with the probabilities resample moved to GPU:
measured ~50-60x on that stage (158s -> ~0.3s), 99.99% voxel agreement vs CPU
(boundary jitter only, no lesion-level change) -> ~1.85x whole pipeline, no accuracy loss.

Includes the fork-specific fix for determine_do_sep_z_and_axis returning the anisotropy
axis as a bare scalar (standard nnU-Net returns a 1-element list), which otherwise
crashes resample_torch's len(axis) assert on the separate-z path.

Wired via the existing EPAI_SCRIPT_PATH hook in _run_epai_inference -- no change to the
dispatch. Enable in prod by setting EPAI_SCRIPT_PATH to scripts/epai_predict.sh after the
end-to-end parity check (seg + findings CSV identical to the bare-CLI path).
---
 flask-server/scripts/epai_predict.py | 108 +++++++++++++++++++++++++++
 flask-server/scripts/epai_predict.sh |   8 ++
 2 files changed, 116 insertions(+)
 create mode 100644 flask-server/scripts/epai_predict.py
 create mode 100644 flask-server/scripts/epai_predict.sh

diff --git a/flask-server/scripts/epai_predict.py b/flask-server/scripts/epai_predict.py
new file mode 100644
index 00000000..f42c0e0b
--- /dev/null
+++ b/flask-server/scripts/epai_predict.py
@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+"""Optimized ePAI inference: GPU-resident export + the fork's own findings CSV.
+
+Drop-in for the bare `nnUNetv2_predict_from_modelfolder` call in
+services/auto_segmentor.py, wired in via the EPAI_SCRIPT_PATH hook (which invokes:
+`bash <script> session_dir case_id input_dir save_dir input_csv output_csv ckpt_path`).
+Runs inside the `epai` conda env, where `import nnunetv2` resolves to qchen76's fork.
+
+Why it exists: profiling showed ePAI's export/convert stage is ~158s of CPU
+resampling of the 26-channel logits back to native resolution -- as expensive as
+inference. Moving that resample to the GPU cuts it to ~0.3s (measured ~50-60x,
+99.99% voxel agreement -- boundary jitter only, no lesion-level change). It reuses
+the fork's own `export_prediction_from_logits`, which does BOTH the segmentation
+export AND the tumor-findings CSV (PDAC/cyst/PNET stats), so the findings report is
+produced identically -- only the resample runs on GPU.
+
+Two fork-specific fixes vs the standard-nnU-Net LesionSegmenter wrapper:
+  1. determine_do_sep_z_and_axis in the fork returns the anisotropy axis as a bare
+     scalar on the separate-z path, so resample_torch's `assert len(axis)==1` crashes
+     (standard nnU-Net returns a 1-element list). Coerce scalar -> [int(axis)].
+  2. Patch ConfigurationManager.resampling_fn_probabilities -> the GPU resampler.
+"""
+import os
+import sys
+from functools import partial
+
+import numpy as np
+import torch
+from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
+from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO
+from nnunetv2.preprocessing.preprocessors.default_preprocessor import DefaultPreprocessor
+from nnunetv2.inference.export_prediction import export_prediction_from_logits
+from nnunetv2.preprocessing.resampling.resample_torch import resample_torch_fornnunet
+import nnunetv2.preprocessing.resampling.resample_torch as _rt
+import nnunetv2.utilities.plans_handling.plans_handler as _ph
+
+
+def _apply_fork_fixes():
+    # (1) coerce the anisotropy axis back to a list so `assert len(axis)==1` holds
+    _orig_determine = _rt.determine_do_sep_z_and_axis
+
+    def _fixed_determine(*a, **k):
+        do_sep, axis = _orig_determine(*a, **k)
+        if do_sep and np.isscalar(axis):
+            axis = [int(axis)]
+        return do_sep, axis
+
+    _rt.determine_do_sep_z_and_axis = _fixed_determine
+
+    # (2) GPU-resident probabilities resampling (the ~158s -> ~0.3s export lever)
+    _ph.ConfigurationManager.resampling_fn_probabilities = property(
+        lambda self: partial(resample_torch_fornnunet, device=torch.device("cuda"), is_seg=False)
+    )
+
+
+def main():
+    # Positional contract from services/auto_segmentor.py's EPAI_SCRIPT_PATH hook.
+    if len(sys.argv) < 8:
+        print("usage: epai_predict.py session_dir case_id input_dir save_dir "
+              "input_csv output_csv ckpt_path", file=sys.stderr)
+        sys.exit(2)
+    _session_dir, case_id, input_dir, save_dir, _input_csv, output_csv, ckpt_path = sys.argv[1:8]
+
+    if not torch.cuda.is_available():
+        print("ERROR: CUDA not available", file=sys.stderr)
+        sys.exit(1)
+
+    _apply_fork_fixes()
+
+    predictor = nnUNetPredictor(
+        tile_step_size=float(os.getenv("EPAI_STEP_SIZE", "0.5")),
+        use_mirroring=os.getenv("EPAI_DISABLE_TTA", "0").strip().lower() not in {"1", "true", "yes", "on"},
+        perform_everything_on_device=True,
+        device=torch.device("cuda"),
+        allow_tqdm=False,
+        verbose=False,
+    )
+    predictor.initialize_from_trained_model_folder(
+        ckpt_path, use_folds=("all",),
+        checkpoint_name=os.getenv("EPAI_CHECKPOINT_NAME", "checkpoint_final.pth"),
+    )
+
+    rw = SimpleITKIO()
+    pp = DefaultPreprocessor(verbose=False)
+    pm, cm = predictor.plans_manager, predictor.configuration_manager
+
+    ct = os.path.join(input_dir, f"{case_id}_0000.nii.gz")
+    if not os.path.exists(ct):
+        print(f"ERROR: input not found: {ct}", file=sys.stderr)
+        sys.exit(1)
+
+    data, props = rw.read_images([ct])
+    dpp, _, pprops = pp.run_case_npy(data, None, props, pm, cm, predictor.dataset_json)
+    logits = predictor.predict_logits_from_preprocessed_data(torch.from_numpy(dpp))
+
+    # The fork's own export: converts logits -> native-res segmentation (GPU resample,
+    # thanks to the patch) AND writes the tumor-findings row into output_csv. Writes the
+    # segmentation to <save_dir>/<case_id>.nii.gz.
+    out_trunc = os.path.join(save_dir, case_id)
+    export_prediction_from_logits(
+        logits, pprops, cm, pm, predictor.dataset_json, out_trunc,
+        save_probabilities=False, output_csv_path=output_csv,
+    )
+    print(f"done {case_id}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/flask-server/scripts/epai_predict.sh b/flask-server/scripts/epai_predict.sh
new file mode 100644
index 00000000..9a3e6b75
--- /dev/null
+++ b/flask-server/scripts/epai_predict.sh
@@ -0,0 +1,8 @@
+#!/bin/bash
+# Wrapper for the EPAI_SCRIPT_PATH hook in services/auto_segmentor.py, invoked as:
+#   bash epai_predict.sh session_dir case_id input_dir save_dir input_csv output_csv ckpt_path
+# The caller has already activated the `epai` conda env, so `python` resolves to the
+# fork's interpreter. Delegates to epai_predict.py (GPU-resident export + findings CSV).
+set -e
+DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+exec python "$DIR/epai_predict.py" "$@"