Overview
Note: This model is an artificial intelligence tool capable of automatically segmenting pan-cancer lesions in 3D computed tomography (CT) scans. As with all medical AI, this model remains an assistive tool: it does not replace the expertise of a radiologist and requires clinical validation for diagnostic use.
This repository is an ONNX export of the vanilla baseline trained by KS987 for the CVPR 2026 — Foundation Models for Pan-cancer Segmentation in CT Images.
- Original weights (PyTorch nnU-Net checkpoint): KS987/CVPR2026-PanCancerSeg-baseline
- Training code, custom trainer & failure analysis: github.com/Kappapapa123/CVPR2026-PanCancerSeg-baseline
No retraining was performed here. This export only converts the released checkpoint_final.pth (fold 0) into a single, self-contained ONNX graph for deployment outside the nnU-Net Python stack (ONNX Runtime / TensorRT).
What this export does (and why)
- Architecture reconstruction: the network is rebuilt from
checkpoint_final.pth. plans.json/dataset.jsonreconstructed from the checkpoint. The original Hugging Face repo does not actually contain these two files (despite what its model card's file listing suggests). Both are embedded by nnU-Net in every checkpoint's so they were extracted directly fromcheckpoint_final.pthrather than reconstructed by hand — guaranteeing an exact match with the weights.- Custom trainer:
nnUNetTrainerWandb2000(from the associated GitHub repo) is required to reload the checkpoint; it was registered into the local nnU-Net installation for this export only and does not affect the exported graph. - Single-file ONNX, opset 17: exported with the legacy (TorchScript-based)
torch.onnxexporter rather than the newerdynamo-based one. The dynamo exporter externalizes weights into a companion.onnx.datafile by default regardless of model size; since this model (~400 MB) is well under the ~2 GB protobuf limit, the legacy exporter keeps everything in a single.onnxfile — simpler to move/version/deploy. - Validated:
onnx.checker.check_modelpasses, and a forward pass was verified end-to-end with ONNX Runtime (CPU).
Exported with in-house tooling.
Input: x, float32, shape [1, 1, 64, 192, 192] (single-channel CT patch, already normalized/resampled per plans.json; axis order Z, Y, X). This matches the 3d_fullres training patch_size [64, 192, 192] from plans.json.
Output: float32, shape [1, 2, 64, 192, 192] — raw logits for 2 classes (0 = background, 1 = tumor), no softmax/argmax applied, no sliding-window aggregation (single patch only — that logic must be implemented by the calling inference engine).
Pre-compiled TensorRT Cache Delivery
Along with the ONNX model weights (unet_pansegx.onnx), this repository provides an optional TensorRT execution cache (trt_engine_cache_fp16.zip).
This cache eliminates the compilation/parsing time (which can take several minutes) during the first inference by loading a graph already optimized and compiled for the target GPU.
The ONNX graph is exported with fully static shapes ([1, 1, 64, 192, 192]), so TensorRT can consume it directly without graph splitting, ensuring maximum performance.
⚠️ The TensorRT cache is tied to the exact ONNX input shape. It must be (re)compiled against the current
[1, 1, 64, 192, 192]graph; any cache built for a different shape will be rejected and ONNX Runtime will recompile a valid engine on first inference.
⚠️ Strict Compatibility Requirements: The TensorRT engine contained in this archive was compiled ahead-of-time. It will nearly only work if the target host environment matches the following configuration:
- GPU Hardware:
sm86(NVIDIA Ampere Architecture). Compatible with: RTX 3070, RTX 3080, RTX 3090, RTX A10G, RTX A40. - Precision:
FP16 - ONNX Runtime: Version
1.19.2 - TensorRT: Version
10.3.0 - CUDA: Version
12.2(or a compatible 12.x series used by TRT 10.3 libraries) - OS: Ubuntu
22.04LTSx86_64
What if my system is not compatible?
If your container crashes (e.g., throwing a TensorRT engine failed to load error), or if you are using a different GPU architecture (e.g., Tesla T4, Ada Lovelace RTX 4x00): Do not use, or simply delete, the extracted files from this zip archive. Run the ONNX inference using the TensorRT backend as usual. ONNX Runtime will detect the missing or mismatched cache and will automatically recompile/re-parse a new, valid engine tailored to your machine. This takes a few minutes on the very first inference, but all subsequent ones will be accelerated.
Intended Use
- Primary use: research use of the ONNX graph within an ONNX Runtime / TensorRT sliding-window inference pipeline.
- Out-of-scope: standalone/clinical use without a proper sliding-window + post-processing pipeline around it; any use outside CT; autonomous diagnosis without clinical review.
Training Details (unchanged from the original checkpoint)
- Dataset: FLARE-MedFM/FLARE-Task1-Pancancer — 10,762 cases (partially labeled
train_label+DeepLesion5K-MedSAM2pseudo-labels). - Framework: nnU-Net v2.6.4, Residual Encoder U-Net Medium (
nnUNetResEncUNetMPlans), configuration3d_fullres. - Trainer:
nnUNetTrainerWandb2000, 2000 epochs, fold 0 (official competition split viasplits_final.json). - Task: binary lesion segmentation (0 = background, 1 = tumor). Only the primary lesion is annotated per training case; metastatic lesions may be present but unlabeled.
Evaluation & Performance (from the original checkpoint's model card — not re-evaluated in ONNX form)
Evaluated on 50 validation-public cases:
| Metric | Mean | Median | Std |
|---|---|---|---|
| Lesion DSC | 0.3278 | 0.2318 | 0.2883 |
| Lesion NSD (2 mm) | 0.1797 | 0.1046 | 0.2029 |
The ONNX export was checked for numerical/behavioral equivalence at the graph level only (checker + a forward pass on random input, shapes match). It was not re-run against the full validation set — the metrics above are inherited from the original PyTorch checkpoint's evaluation and assumed to transfer, since the export changes neither weights nor architecture (only deep supervision is removed, which does not affect the final head).
Limitations
- Inherited from the original model: partial labels (dominant failure mode is scattered false-positive fragmentation — see the failure analysis).
- Not re-validated end-to-end in ONNX form against the reference Dice/NSD metrics (see above).
Citation
If you use this export, please cite this repository, the original model, and the underlying framework/dataset:
This ONNX export:
EOLP. (2026). PANSEGX: An ONNX export of the vanilla baseline for CVPR 2026 Pan-Cancer Segmentation. Hugging Face. https://huggingface.co/EOLP/PANSEGX
Model checkpoint:
KS987. CVPR2026-PanCancerSeg-baseline: Vanilla nnU-Net v2 Baseline for CVPR 2026 Task 1: Pan-Cancer Segmentation. Hugging Face. https://huggingface.co/KS987/CVPR2026-PanCancerSeg-baseline
nnU-Net Framework:
Isensee, F., Jaeger, P.F., Kohl, S.A.A. et al. (2021). nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature Methods 18, 203–211.
FLARE-Task1-Pancancer Dataset:
FLARE-MedFM/FLARE-Task1-Pancancer. Hugging Face Datasets. https://huggingface.co/datasets/FLARE-MedFM/FLARE-Task1-Pancancer
Disclaimer: This model is provided for research and assistive purposes only. It does not replace medical advice. This ONNX export is an unofficial derivative of KS987's checkpoint; please refer to the original repository for the authoritative model card and license terms.