""" Vbai-2.6AD — configuration. Paired multimodal early-Alzheimer's detection: a 3D MRI volume plus a panel of 13 biomarkers, fused into one representation. -------------------------------------------------------------------------- YOU MUST SET YOUR OWN PATHS. -------------------------------------------------------------------------- No data location is hard-coded. Point the environment variables below at your own files before running anything: VBAI_DATASET_ROOT root of your imaging + tabular data VBAI_VOLUME_ROOT root of the volume files referenced by the manifest VBAI_MODEL_SAVE_ROOT where checkpoints are written If a variable is unset, the loader walks up the project tree looking for a `Datasets` directory. If that fails too, the data-preparation step raises an explicit error rather than guessing. Nothing here describes or names a particular corpus. Bring your own data; the expected column contract is FEATURE_NAMES below. """ import os from dataclasses import dataclass, field from typing import List PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__)) def _walk_up_for_marker(start: str, marker_subdirs=("Datasets",), max_levels: int = 6): """Walk upwards from `start` until a directory containing all markers is found.""" cur = os.path.abspath(start) for _ in range(max_levels): if all(os.path.isdir(os.path.join(cur, m)) for m in marker_subdirs): return cur parent = os.path.dirname(cur) if parent == cur: break cur = parent return None def _resolve_dataset_root() -> str: """Resolution order: environment variable → project-tree search → default.""" env = os.environ.get("VBAI_DATASET_ROOT") if env and os.path.isdir(env): return env walked = _walk_up_for_marker(PROJECT_ROOT, marker_subdirs=("Datasets",)) if walked is not None: return os.path.join(walked, "Datasets") # Fall through to a relative default; data preparation will report clearly # if nothing is there. SET VBAI_DATASET_ROOT TO YOUR OWN PATH. return os.path.normpath(os.path.join(PROJECT_ROOT, "..", "Datasets")) def _resolve_model_save_root() -> str: env = os.environ.get("VBAI_MODEL_SAVE_ROOT") if env: return env walked = _walk_up_for_marker(PROJECT_ROOT, marker_subdirs=("Datasets",)) if walked is not None: return os.path.join(walked, "Models", "Vbai-2.6AD") return os.path.normpath(os.path.join(PROJECT_ROOT, "..", "Models", "Vbai-2.6AD")) DATASET_ROOT = _resolve_dataset_root() MODEL_SAVE_ROOT = _resolve_model_save_root() # Root of the volume files. The visit manifest stores relative paths; this is # what they are resolved against. SET VBAI_VOLUME_ROOT TO YOUR OWN PATH. VOLUME_ROOT = os.environ.get("VBAI_VOLUME_ROOT") or os.path.join(DATASET_ROOT, "volumes") # Kept for backward compatibility with scripts that expect these names. TBM_ROOT = VOLUME_ROOT TBM_CSV = os.environ.get("VBAI_VOLUME_MANIFEST") or os.path.join(VOLUME_ROOT, "manifest.csv") # Which volume modality is in use. Selected by the extraction scripts through # the --tbm / --t1 flag, which sets this variable before config is imported. # A checkpoint trained on one modality must never be fed the other. USE_TBM = bool(int(os.environ.get("VBAI_USE_TBM", "0"))) CACHE_DIR = os.path.join(PROJECT_ROOT, "_cache") os.makedirs(CACHE_DIR, exist_ok=True) PAIRED_PARQUET_T1 = os.path.join(CACHE_DIR, "paired_visits.parquet") PAIRED_PARQUET_TBM = os.path.join(CACHE_DIR, "paired_visits_tbm.parquet") PAIRED_PARQUET = PAIRED_PARQUET_TBM if USE_TBM else PAIRED_PARQUET_T1 # Tabular feature order — a fixed contract relied on everywhere downstream. # Your table must provide these columns (missing values are allowed and are # handled explicitly through a per-feature mask; see NUM_TABULAR_INPUTS). FEATURE_NAMES: List[str] = [ "Age", # demographic "Sex", # 0 = F, 1 = M "MMSE", # cognitive "CDRSB", # cognitive (CDR sum of boxes) "APOE4_count", # genetic, 0/1/2 e4 alleles "CSF_ABETA42", # CSF "CSF_TAU", # CSF "CSF_PTAU", # CSF "CSF_AB42_AB40", # CSF ratio "PLASMA_PTAU", # blood "PLASMA_NFL", # blood "PLASMA_AB42_AB40", # blood ratio "PLASMA_GFAP", # blood ] NUM_FEATURES = len(FEATURE_NAMES) # 13 # One value plus one missing-mask bit per feature. The mask is not decoration: # an unmeasured biomarker must stay distinguishable from a normal one. NUM_TABULAR_INPUTS = NUM_FEATURES * 2 CLASS_NAMES = ["CN", "MCI", "AD"] DIAGNOSIS_MAP = {"CN": 0, "MCI": 1, "Dementia": 2, "AD": 2, "EMCI": 1, "LMCI": 1, "SMC": 0} @dataclass class ModelConfig: mri_input_shape: tuple = (1, 96, 96, 96) mri_encoder_channels: List[int] = field(default_factory=lambda: [32, 64, 128, 256]) mri_bottleneck_channels: int = 512 mri_feature_dim: int = 512 mri_dropout: float = 0.4 use_cbam: bool = True use_se_block: bool = True num_tabular_inputs: int = NUM_TABULAR_INPUTS tabular_hidden_dims: List[int] = field(default_factory=lambda: [128, 256]) tabular_feature_dim: int = 256 tabular_dropout: float = 0.3 fusion_dim: int = 512 fusion_num_heads: int = 8 fusion_dropout: float = 0.3 num_classes: int = 3 progression_hidden_dim: int = 256 max_progression_months: int = 120 num_time_bins: int = 24 # Modality dropout during training: teaches the model to survive a missing # arm at inference instead of collapsing. p_drop_mri: float = 0.15 p_drop_tab: float = 0.15 # Feature-wise random masking, simulating biomarkers absent at inference. p_feature_mask: float = 0.20 @dataclass class TrainingConfig: seed: int = 42 device: str = "cuda" num_workers: int = 4 pin_memory: bool = True mixed_precision: bool = True # Phase 1 — MRI encoder pretraining phase1_epochs: int = 40 phase1_batch_size: int = 4 phase1_lr: float = 3e-4 phase1_weight_decay: float = 1e-4 # Phase 2 — tabular encoder pretraining phase2_epochs: int = 60 phase2_batch_size: int = 64 phase2_lr: float = 1e-3 phase2_weight_decay: float = 1e-4 # Phase 3 — joint fusion on paired visits phase3_epochs: int = 40 phase3_batch_size: int = 4 phase3_lr_backbone: float = 1e-5 phase3_lr_fusion: float = 5e-4 phase3_weight_decay: float = 1e-4 # Loss weights w_cls_fused: float = 1.0 w_cls_mri: float = 0.3 w_cls_tab: float = 0.3 w_prog: float = 0.5 w_contrastive: float = 0.2 focal_gamma: float = 1.0 label_smoothing: float = 0.05 grad_clip: float = 1.0 val_split: float = 0.15 test_split: float = 0.15 # subject-level holdout, never visit-level early_stopping_patience: int = 20 min_epochs_before_es: int = 25 save_dir: str = MODEL_SAVE_ROOT @dataclass class DataConfig: nifti_target_shape: tuple = (96, 96, 96) pair_window_months: int = 6 # MRI ↔ biomarker date tolerance progression_horizon_months: int = 60 # 5-year look-ahead for MCI → AD aug_rotation_range: float = 8.0 aug_flip_prob: float = 0.5 aug_noise_std: float = 0.02 aug_gamma_range: tuple = (0.85, 1.15) # Optional hippocampus-focused crop: the brain bounding box is found, then # a centre crop is taken at these ratios and resized to nifti_target_shape. # x: left-right (both hemispheres), y: anterior-posterior, z: inferior-superior hippocampus_crop_enabled: bool = False hippo_x_range: tuple = (0.10, 0.90) hippo_y_range: tuple = (0.25, 0.70) hippo_z_range: tuple = (0.15, 0.65)