Dianyuan/VDPA

# VDPA: Variance-Driven Dual-Path Attention for CT-based KRAS Prediction

This repository provides the PyTorch implementation and checkpoint for VDPA, a variance-driven dual-path attention network for CT-based prediction of KRAS mutation status in non-small cell lung cancer.

Overview

VDPA integrates local sparse attention and learnable global aggregation through a variance-driven fusion mechanism. The released model uses three tumor-containing CT slices as input and produces a patient-level KRAS mutation probability.

Core configuration:

Item	Setting
Input	3 x 512 x 512 CT slices
Patch size	16
Embedding dimension	768
Transformer depth	10
Attention heads	12
Class token	Yes
Number of classes	2
Parameters	77.6M

Repository Structure

src/vdpa/model.py                  VDPA model definition
src/vdpa/commn.py                  Attention and pooling modules
scripts/check_checkpoint_compatibility.py
scripts/evaluate_npy_checkpoint.py
scripts/prepare_tcia_lobe3_npy.py
scripts/plot_metrics_and_calibration.py
configs/vdpa_checkpoint_config.json
checkpoints/vdpa_institutional_best_model.pth
results/

Installation

conda create -n vdpa-kras python=3.9 -y
conda activate vdpa-kras
pip install -r requirements.txt

Install a CUDA-compatible PyTorch build if GPU inference is required.

Checkpoint Compatibility

Run the following command after downloading the repository:

python scripts/check_checkpoint_compatibility.py \
  --checkpoint checkpoints/vdpa_institutional_best_model.pth

Expected output:

missing: []
unexpected: []
parameter_count: 77595168

Data Format

The evaluation script expects NPY files arranged by class and patient:

data/KRAS_NPY/val/KRAS/<patient_id>/<case>.npy
data/KRAS_NPY/val/mut-KRAS/<patient_id>/<case>.npy

Each .npy file should contain one three-slice CT input in [3, H, W] format. [H, W, 3] arrays are also supported and are transposed automatically.

Class labels:

Folder	Label
KRAS	0
mut-KRAS	1

Evaluation

python scripts/evaluate_npy_checkpoint.py \
  --data-root data/KRAS_NPY \
  --split val \
  --checkpoint checkpoints/vdpa_institutional_best_model.pth \
  --out-dir results/checkpoint_eval \
  --batch-size 2 \
  --threshold 0.5

Outputs:

results/checkpoint_eval/image_predictions.csv
results/checkpoint_eval/patient_predictions.csv
results/checkpoint_eval/metrics.json

Patient-level probability is computed by averaging mutation probabilities across all inputs belonging to the same patient.

TCIA Preprocessing Helper

The repository includes a helper script for preparing three-slice NPY inputs from the public TCIA NSCLC Radiogenomics collection:

python scripts/prepare_tcia_lobe3_npy.py \
  --tcia-root data/TCIA_NSCLC_Radiogenomics_KRAS \
  --out-root data/TCIA_KRAS_lobe3_npy

The script reads TCIA metadata, selects CT series, extracts three axial CT slices according to the recorded tumor lobe zone, and writes NPY inputs using the same class-folder format as the evaluation script.

Checkpoints and Git LFS

The model checkpoint is larger than GitHub's standard 100 MB file limit. Use Git LFS before committing the checkpoint:

git lfs install
git lfs track "*.pth"
git add .gitattributes

Data Availability

Institutional CT data are not included in this repository. Public TCIA data should be obtained from The Cancer Imaging Archive according to its data access policy.

Citation

@article{vdpa_kras_nsclc,
  title = {Variance-Driven Dual-Path Attention for CT Prediction of KRAS Mutation Status in Non-Small Cell Lung Cancer},
  author = {Fu, Guobin and colleagues},
  year = {2026}
}

License

This code is released under the MIT License.