""" Vbai-2.6AD Datasets =================== Reads the cached visit manifest at _cache/paired_visits.parquet, which your own data-preparation step must produce. One row per visit: a volume path, the 13 biomarker columns with their per-feature masks, the label and the progression fields. Three dataset modes: * mode="mri" → MRI + label (Phase 1 pretrain) * mode="tab" → tabular features + label (Phase 2 pretrain) * mode="multi" → MRI + tabular + label + progression (Phase 3 fusion) Tabular feature contract: 2 * NUM_FEATURES floats per sample. [normalized values..., missing-mask bits...] A feature with missing-mask=0 has its value zeroed (after normalization). """ from __future__ import annotations import os import random import numpy as np import pandas as pd import torch from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler from scipy.ndimage import zoom, rotate import config as C try: import nibabel as nib HAS_NIBABEL = True except Exception: HAS_NIBABEL = False # Optional pre-decoded .npy cache, searched in order. Decoding NIfTI is the # slowest part of an epoch, so a local-disk cache pays for itself quickly. # YOU MUST SET YOUR OWN PATH: point VBAI_NPY_CACHE at a fast local directory, # or leave it unset to use the in-project cache. _NPY_CACHE_DIRS = [d for d in [ os.environ.get("VBAI_NPY_CACHE"), os.path.join(C.PROJECT_ROOT, "_cache", "volume_npy"), ] if d] def _try_load_cached(image_id: str) -> np.ndarray | None: if not image_id: return None for d in _NPY_CACHE_DIRS: p = os.path.join(d, f"{image_id}.npy") if os.path.exists(p): try: return np.load(p).astype(np.float32) except Exception: return None return None # ---------- MRI loading + augmentation ---------- def _hippocampus_crop(data: np.ndarray, dcfg: C.DataConfig) -> np.ndarray: """Find the brain bounding box, then crop to a hippocampus-focused sub-region.""" mask = data > 0 if not mask.any(): return data coords = np.argwhere(mask) mn = coords.min(axis=0); mx = coords.max(axis=0) size = mx - mn + 1 rx, ry, rz = dcfg.hippo_x_range, dcfg.hippo_y_range, dcfg.hippo_z_range x0, x1 = int(mn[0] + size[0] * rx[0]), int(mn[0] + size[0] * rx[1]) y0, y1 = int(mn[1] + size[1] * ry[0]), int(mn[1] + size[1] * ry[1]) z0, z1 = int(mn[2] + size[2] * rz[0]), int(mn[2] + size[2] * rz[1]) cropped = data[x0:x1+1, y0:y1+1, z0:z1+1] return cropped def _load_nifti(path: str, target_shape=(96, 96, 96), hippocampus_crop: bool = False, dcfg: C.DataConfig = None) -> np.ndarray: img = nib.load(path) data = img.get_fdata().astype(np.float32) if data.ndim == 4: data = data[..., 0] mask = data > 0 if mask.sum() > 0: vals = data[mask] lo, hi = np.percentile(vals, [1.0, 99.0]) data = np.clip(data, lo, hi) m, s = vals.mean(), vals.std() if s > 0: data = (data - m) / s data[~mask] = 0 # Hippocampus-focused crop: narrows the content, which raises effective resolution if hippocampus_crop: data = _hippocampus_crop(data, dcfg or C.DataConfig()) if data.shape != target_shape: f = [t / s for t, s in zip(target_shape, data.shape)] data = zoom(data, f, order=1) return data.astype(np.float32) class MRIAugment3D: def __init__(self, dcfg: C.DataConfig): self.cfg = dcfg def __call__(self, vol: np.ndarray) -> np.ndarray: if random.random() < 0.5: angle = random.uniform(-self.cfg.aug_rotation_range, self.cfg.aug_rotation_range) axes = random.choice([(0, 1), (0, 2), (1, 2)]) vol = rotate(vol, angle, axes=axes, reshape=False, order=1, mode="nearest") for ax in range(3): if random.random() < self.cfg.aug_flip_prob: vol = np.flip(vol, axis=ax).copy() if random.random() < 0.3: vol = vol + np.random.normal(0, self.cfg.aug_noise_std, vol.shape).astype(np.float32) if random.random() < 0.3: g = random.uniform(*self.cfg.aug_gamma_range) mn = vol.min(); rg = vol.max() - mn if rg > 0: vol = ((vol - mn) / rg) ** g * rg + mn return vol.astype(np.float32) # ---------- Tabular normalization ---------- class TabularNormalizer: """Robust z-score on observed (non-missing) values per feature, fit on training set.""" def __init__(self): self.mean: np.ndarray | None = None self.std: np.ndarray | None = None def fit(self, df: pd.DataFrame): means, stds = [], [] for f in C.FEATURE_NAMES: v = pd.to_numeric(df[f], errors="coerce").dropna().values.astype(np.float64) if len(v) > 1: m = float(np.median(v)) s = float(np.median(np.abs(v - m)) * 1.4826) # MAD → std if s < 1e-8: s = float(v.std()) if v.std() > 1e-8 else 1.0 else: m, s = 0.0, 1.0 means.append(m); stds.append(s) self.mean = np.asarray(means, dtype=np.float32) self.std = np.asarray(stds, dtype=np.float32) def transform(self, values: np.ndarray, mask: np.ndarray) -> np.ndarray: z = (values - self.mean) / self.std z = np.where(mask > 0.5, z, 0.0) # zero out missing return np.concatenate([z, mask.astype(np.float32)], axis=-1) def state_dict(self): return {"mean": self.mean.tolist() if self.mean is not None else None, "std": self.std.tolist() if self.std is not None else None} def load_state_dict(self, sd): self.mean = np.asarray(sd["mean"], dtype=np.float32) self.std = np.asarray(sd["std"], dtype=np.float32) # ---------- Subject-level split (no leakage between train/val/test) ---------- def subject_split(df: pd.DataFrame, val_frac=0.15, test_frac=0.15, seed=42): rng = np.random.RandomState(seed) ptids = np.array(sorted(df["ptid"].unique())) rng.shuffle(ptids) n = len(ptids) n_test = int(round(n * test_frac)) n_val = int(round(n * val_frac)) test_ids = set(ptids[:n_test]) val_ids = set(ptids[n_test:n_test + n_val]) train_ids = set(ptids[n_test + n_val:]) print(f"[split] subjects → train {len(train_ids)} / val {len(val_ids)} / test {len(test_ids)}") return train_ids, val_ids, test_ids # ---------- Core paired dataset ---------- class PairedVisitDataset(Dataset): """ One sample = one MRI scan with paired biomarkers + (optional) progression labels. Setting mode controls which fields are loaded: "mri" — only mri + label (skips biomarker columns) "tab" — only biomarkers + label (skips MRI loading) "multi" — both """ def __init__(self, df: pd.DataFrame, normalizer: TabularNormalizer, mode: str = "multi", augment: bool = False, dcfg: C.DataConfig = None, mcfg: C.ModelConfig = None, train_modality_dropout: bool = False): self.df = df.reset_index(drop=True).copy() self.norm = normalizer self.mode = mode self.augment = augment self.dcfg = dcfg or C.DataConfig() self.mcfg = mcfg or C.ModelConfig() self.augmenter = MRIAugment3D(self.dcfg) if augment else None self.modality_dropout = train_modality_dropout def __len__(self): return len(self.df) def _get_tab(self, row, training: bool): vals = np.array([row[f] for f in C.FEATURE_NAMES], dtype=np.float32) mask = np.array([row[f"feat_mask_{f}"] for f in C.FEATURE_NAMES], dtype=np.float32) # NaN safety vals = np.where(np.isnan(vals), 0.0, vals) # Stochastic feature masking during training (simulate missing inputs) if training and self.mcfg.p_feature_mask > 0: drop = np.random.rand(len(C.FEATURE_NAMES)) < self.mcfg.p_feature_mask mask = np.where(drop, 0.0, mask) return self.norm.transform(vals, mask).astype(np.float32) def _get_mri(self, row): # Fast path: pre-decoded .npy on local disk vol = _try_load_cached(row.get("image_id")) if vol is None: vol = _load_nifti(row["nifti_path"], self.dcfg.nifti_target_shape) if self.augment and self.augmenter: vol = self.augmenter(vol) return torch.from_numpy(np.ascontiguousarray(vol)).unsqueeze(0).float() def __getitem__(self, idx): row = self.df.iloc[idx] out = { "label": torch.tensor(int(row["label"]), dtype=torch.long), "has_progression": torch.tensor(bool(row["has_progression"]), dtype=torch.bool), "will_progress": torch.tensor(float(row["will_progress"]), dtype=torch.float32), "progression_months": torch.tensor(float(row["months_to_conversion"]), dtype=torch.float32), "ptid": str(row["ptid"]), } load_mri = self.mode in ("mri", "multi") load_tab = self.mode in ("tab", "multi") # Modality dropout (Phase 3 only) if self.modality_dropout and self.mode == "multi": r = random.random() if r < self.mcfg.p_drop_mri: load_mri = False elif r < self.mcfg.p_drop_mri + self.mcfg.p_drop_tab: load_tab = False if load_mri: out["mri"] = self._get_mri(row) if load_tab: out["tab"] = torch.from_numpy(self._get_tab(row, training=self.augment)) out["has_mri"] = torch.tensor(load_mri, dtype=torch.bool) out["has_tab"] = torch.tensor(load_tab, dtype=torch.bool) return out def collate_pad(batch): """Collate that handles optional mri/tab tensors per-sample.""" keys = ["label", "has_progression", "will_progress", "progression_months", "has_mri", "has_tab"] out = {k: torch.stack([b[k] for b in batch]) for k in keys} # MRI: only stack if all present (modality dropout makes mixed batches rare in practice; # we drop unmatched samples to None at batch level to keep things simple) if all("mri" in b for b in batch): out["mri"] = torch.stack([b["mri"] for b in batch]) if all("tab" in b for b in batch): out["tab"] = torch.stack([b["tab"] for b in batch]) out["ptid"] = [b["ptid"] for b in batch] return out # ---------- Helpers ---------- def get_class_weights(labels: np.ndarray, num_classes: int = 3) -> torch.Tensor: counts = np.bincount(labels, minlength=num_classes).astype(np.float32) counts[counts == 0] = 1.0 w = 1.0 / counts w = w / w.sum() * num_classes return torch.tensor(w, dtype=torch.float32) def get_weighted_sampler(labels: np.ndarray) -> WeightedRandomSampler: counts = np.bincount(labels) sw = 1.0 / counts[labels] return WeightedRandomSampler(torch.from_numpy(sw).float(), len(sw), replacement=True) def load_paired() -> pd.DataFrame: if not os.path.exists(C.PAIRED_PARQUET): raise FileNotFoundError( f"Visit manifest not found: {C.PAIRED_PARQUET}\n" "Build it from your own data first. YOU MUST SET YOUR OWN PATHS " "(see config.py)." ) return pd.read_parquet(C.PAIRED_PARQUET)