PR3DICTR: A modular AI framework for medical 3D image-based detection and outcome prediction

1. Problem Statement

The field of medical imaging is increasingly relying on 3D multi-modal data (CT, PET, MRI) and deep learning (DL) for diagnosis and outcome prediction. However, the development of DL models in this domain faces several critical challenges:

• Lack of Standardization: Research groups often develop diverse, non-standardized workflows, leading to inconsistencies in data handling, training, and evaluation.
• High Developmental Burden: Building models from scratch requires substantial computational resources and technical expertise, often resulting in redundant code re-implementation across different projects.
• Framework Limitations: Existing solutions are often too broad (e.g., general PyTorch) to be immediately applicable to specific medical tasks, or too narrow (e.g., no-code tools) to allow necessary customization for complex medical outcomes.
• Data Complexity: Medical data involves varied modalities, missing values, class imbalances, and the need to integrate 3D volumetric images with tabular clinical data, which complicates preprocessing and model architecture design.

2. Methodology

The authors introduce PR3DICTR (Platform for Research in 3D Image Classification and sTandardised tRaining), a modular framework built on PyTorch and MONAI. The methodology is structured around a six-step pipeline:

A. Data Curation and Preprocessing

• Input Requirements: The framework assumes data is pre-curated into a standardized format. It requires:
  • A CSV file containing patient IDs, split labels (train/val/test), and target endpoints.
  • 3D volumetric data stored as NumPy (.npy) files with consistent dimensions.
  • A specific directory structure organized by PatientID.
• Preprocessing: The framework handles value clipping, normalization, and cropping. It supports deterministic transforms (applied to all splits) and non-deterministic transforms (random flips, rotations, MixUp) applied only to the training set.

B. Modular Architecture

The framework is designed with a "plug-and-play" philosophy, allowing users to swap components without rewriting the entire pipeline.

• Configuration-Driven: A central YAML configuration file defines all aspects of the experiment (hyperparameters, architecture, augmentation, loss functions). Users can start a project with as little as two lines of code (load config, run experiment).
• Model Structure:
  • Image Encoder: Supports various backbones (ResNet, DenseNet, EfficientNetV2, ConvNeXt, ViT, TransRP) to extract features from 3D volumes. Multiple modalities (e.g., CT + PET) are stacked along the channel dimension.
  • Output Module: Integrates image features with tabular clinical data via fully connected layers or Vision Transformers (ViT). It supports multi-label classification with independent output heads.
  • MLP Mode: If no images are provided, the framework defaults to a Multi-Layer Perceptron (MLP) for tabular-only data.

C. Training and Optimization

• Cross-Validation: Supports K-fold cross-validation with stratified sampling based on clinical variables.
• Hyperparameter Optimization: Integrates Optuna for automated, efficient hyperparameter search, replacing extensive grid searches.
• Data Loading: Offers four MONAI dataset interfaces (Standard, Cache, SmartCache, Persistent) to balance memory usage and training speed.
• Loss Functions: Includes specialized losses for imbalanced data (Focal, Hill, ASL) and survival analysis (Negative Log-Likelihood).

D. Evaluation and Reproducibility

• Metrics: Automatically computes a wide range of metrics including AUC, Accuracy, F1-score, C-index (for survival), and calibration errors (ACE, ECE, MCE).
• Visualization: Generates ROC curves, confusion matrices, calibration plots, and Kaplan-Meier curves.
• Tracking: Integrates with Weights & Biases (W&B) for real-time experiment tracking.
• Post-hoc Evaluation: Includes a dedicated function to evaluate saved models on external test sets using the same standardized metrics.

3. Key Contributions

1. Unified Framework: PR3DICTR bridges the gap between flexible low-level frameworks (PyTorch) and rigid high-level tools, offering a balanced, modular solution specifically for 3D medical imaging.
2. Standardization: It enforces a consistent pipeline for data loading, training, and evaluation, reducing implementation variability across different research projects (e.g., head and neck vs. lung cancer).
3. Modularity and Flexibility: Users can easily swap encoders, loss functions, and augmentation strategies via configuration files without altering core code. It supports multi-modal fusion (images + tabular data) and various classification tasks (binary, event-based/survival).
4. Reproducibility: By saving configuration files (YAML) and model weights for every experiment, the framework ensures that results can be exactly replicated by collaborators or external reviewers.
5. Accessibility: The "two-line code" entry point significantly lowers the barrier to entry for researchers, allowing them to focus on methodology rather than boilerplate code.

4. Results

• Proof of Concept: The framework was tested on the NSCLC-Radiomics dataset (The Cancer Imaging Archive) for a sex classification task.
• Performance: Using a default ResNet-10 architecture, the model achieved nearly perfect distinction between sexes in the test set with decent calibration.
• Efficiency: The framework successfully streamlined the development of complex models (e.g., NTCP models for xerostomia and dysphagia) by allowing the reuse of core training and evaluation code across different disease sites, requiring only config file modifications.

5. Significance

• Accelerating Research: PR3DICTR reduces the time required to develop and validate new DL models, allowing researchers to rapidly prototype and compare different architectures and strategies.
• Enhancing Clinical Translation: By standardizing workflows and improving reproducibility, the framework fosters trust in AI models, a prerequisite for clinical adoption.
• Future-Proofing: The modular design allows for easy integration of emerging technologies, such as uncertainty quantification, attention maps for interpretability, and multi-class classification, without overhauling the system.
• Community Impact: As an open-source tool, it provides a common language and infrastructure for the medical AI community, facilitating collaboration and meta-analyses across institutions.

In conclusion, PR3DICTR represents a significant step forward in operationalizing deep learning for 3D medical imaging, moving the field from fragmented, custom scripts toward a standardized, modular, and reproducible ecosystem.