Add coreg/atlas layout-2 support and rare-pathology rebalanced sampling

Dataloader now picks the image subfolder from --space when scanning the raw
HuggingFace download layout: native_space -> img/, coreg_space -> coreg_img/,
atlas_space -> atlas_img/. Layout 1 (<uid>/<space>/img/) is unchanged.

Adds optional inverse-prevalence weighted sampling for contrastive
pretraining: pass --pathology_labels_csv together with --rebalance_strategy
(inverse_freq | sqrt_inverse_freq | max_inverse_freq) to switch the
DataLoader to a WeightedRandomSampler that upsamples subjects with rare
positive pathologies. Unlabeled / all-negative subjects keep
--rebalance_base_weight (default 1.0). 8 new unit tests cover the weight
formulas and sampler integration; 5 new tests cover the coreg/atlas layout
discovery.

Author: sezginerr

contrastive-pretraining/README.md

@@ -146,21 +146,32 @@ data_folder/
 ```
 
 **Layout 2 — Batch-based** (HuggingFace dataset format):
+
+The image subdirectory name is chosen automatically from `--space`:
+
+| Space (`--space`) | Image subfolder | HuggingFace repo |
+|-------------------|-----------------|------------------|
+| `native_space` *(default)* | `img/` | [Forithmus/MR-RATE](https://huggingface.co/datasets/Forithmus/MR-RATE) |
+| `coreg_space` | `coreg_img/` | [Forithmus/MR-RATE-coreg](https://huggingface.co/datasets/Forithmus/MR-RATE-coreg) |
+| `atlas_space` | `atlas_img/` | [Forithmus/MR-RATE-atlas](https://huggingface.co/datasets/Forithmus/MR-RATE-atlas) |
+
 ```
-data_folder/
+data_folder/                       # e.g. ./data/MR-RATE-coreg/mri
 ├── batch00/
 │   ├── SUBJECT_001/
-│   │   └── img/
+│   │   └── coreg_img/             # img/ for native, atlas_img/ for atlas
 │   │       ├── SUBJECT_001_t1w-raw-axi.nii.gz
 │   │       ├── SUBJECT_001_flair-raw-axi.nii.gz
 │   │       └── ...
 │   ├── SUBJECT_002/
-│   │   └── img/
+│   │   └── coreg_img/
 │   │       └── ...
 ├── batch01/
 │   └── ...
 ```
 
+After `merge_downloaded_repos.py` consolidates native + coreg + atlas under a single tree, all three subfolders (`img/`, `coreg_img/`, `atlas_img/`) live side-by-side under each study, and `--space` selects which one to load.
+
 Reports are stored in a JSONL file with one entry per subject:
 
 ```json
@@ -230,8 +241,11 @@ FUSION_MODE=late_attn POOLING_STRATEGY=cross_attn NUM_TRAIN_STEPS=50000 \
 | `--pooling_strategy` | `simple_attn`, `cross_attn`, `gated` | `simple_attn` | Volume pooling (used with `late_attn`) |
 | `--data_folder` | — | required | Path to MR data folder |
 | `--jsonl_file` | — | required | Path to reports JSONL file |
-| `--space` | — | `native_space` | Image space subfolder (only for space-based layout) |
+| `--space` | `native_space`, `coreg_space`, `atlas_space` | `native_space` | Selects the image subfolder under each study: `<space>/img/` (layout 1) or `img/` / `coreg_img/` / `atlas_img/` (layout 2) |
 | `--normalizer` | `zscore`, `percentile`, `minmax` | `zscore` | Volume normalization method |
+| `--pathology_labels_csv` | — | — | Path to pathology labels CSV (e.g. `mrrate_labels.csv`). Required for rebalancing. |
+| `--rebalance_strategy` | `inverse_freq`, `sqrt_inverse_freq`, `max_inverse_freq` | — (uniform) | Per-subject sampling weight strategy for upsampling rare pathologies. |
+| `--rebalance_base_weight` | — | `1.0` | Base sampling weight for all-negative / unlabeled subjects. |
 | `--num_train_steps` | — | `100001` | Total training steps |
 | `--lr` | — | `1e-5` | Learning rate |
 | `--results_folder` | — | `./mr_rate_results` | Checkpoint output directory |
@@ -249,6 +263,61 @@ FUSION_MODE=late_attn POOLING_STRATEGY=cross_attn NUM_TRAIN_STEPS=50000 \
 - **W&B integration**: `--wandb` logs loss, learning rate, volume count per step; run ID is persisted for seamless resume across SLURM jobs
 - **Gradient checkpointing**: Nested checkpointing at both volume level (in MRRATE) and chunk level (in sliding encoders) for maximum memory efficiency
 - **Dual checkpoints**: Saves both model-only `.pt` (for inference) and full `.pt` (model + optimizer + scheduler + step, for resume)
+- **Rare-pathology rebalancing** (see below): inverse-prevalence weighted sampling so contrastive batches see uncommon pathologies more often
+
+### Rare-Pathology Rebalancing
+
+Pathology prevalence in MR-RATE is heavily skewed (~43% of studies are all-negative; the rarest pathology, *Hemangioma of vertebral column*, occurs in ~0.23% of studies; *Gliosis* in ~37%). Under uniform sampling, contrastive batches are dominated by common findings and the model rarely sees rare ones. Passing `--pathology_labels_csv` together with `--rebalance_strategy` switches the training DataLoader to a `WeightedRandomSampler` whose per-subject weights are derived from inverse prevalence.
+
+#### How the weighting works
+
+Each subject's report is **multi-label** — a single scan can have zero, one, or many positive pathologies. Every subject is reduced to a single sampling weight `w_i`. Bigger weight = drawn more often. The three strategies differ in how a subject's binary label vector is collapsed into that one number.
+
+Concrete walk-through with two pathologies and three subjects (using real prevalences from `mrrate_labels.csv`: *Hemangioma of vertebral column* ≈ 0.23% → inv-freq ≈ 433; *Gliosis* ≈ 37% → inv-freq ≈ 2.7):
+
+| Subject | Hemangioma (rare) | Gliosis (common) | `inverse_freq` weight |
+|---------|-------------------|------------------|-----------------------|
+| A       | 1                 | 1                | `1 + 433 + 2.7 ≈ 436.7` |
+| B       | 0                 | 1                | `1 +   0 + 2.7 ≈ 3.7`   |
+| C       | 0                 | 0                | `1` (base)              |
+
+Subject A is drawn ~436× more often than C, and ~118× more often than B.
+
+#### Strategy choice
+
+For each pathology `p`, `prevalence[p]` is the positive rate across labeled subjects in the active split, and `inv_freq[p] = 1 / max(prevalence[p], eps)`. The three strategies are different ways to aggregate across a subject's positive labels:
+
+| Strategy | Formula | Subject A's weight | Behavior |
+|----------|---------|--------------------|----------|
+| `inverse_freq` | `base + Σ_p y_p · inv_freq[p]` | `436.7` | **Sum** — rewards multiple co-occurring rare findings. On `mrrate_labels.csv`, rarest pathology mass: 0.23% → ~4.0% (~17× boost). |
+| `sqrt_inverse_freq` | `base + Σ_p y_p · √inv_freq[p]` | `22.4` | Sum with diminishing returns. Rarest pathology mass: 0.23% → ~1.5%. Use when `inverse_freq` over-amplifies. |
+| `max_inverse_freq` | `max(base, max_p y_p · inv_freq[p])` | `433` | **Max** — a subject's weight is set by its single rarest positive pathology. Doesn't matter if they have one rare finding or five — caps combinatorial blow-up. |
+
+Recommendation: start with `inverse_freq`. Switch to `sqrt_inverse_freq` if training loss gets noisy from over-sampling the same handful of rare subjects.
+
+Subjects missing from the labels CSV — and all-negative subjects — receive `--rebalance_base_weight` (default `1.0`), so they remain in the pool but do not dominate it.
+
+#### Enabling it
+
+Add two flags to your existing training command:
+
+```bash
+accelerate launch --multi_gpu --num_processes 4 scripts/run_train.py \
+    --encoder vjepa2 --fusion_mode late \
+    --data_folder /path/to/data \
+    --jsonl_file  /path/to/reports.jsonl \
+    --pathology_labels_csv /path/to/mrrate_labels.csv \
+    --rebalance_strategy inverse_freq
+```
+
+Without those flags, nothing changes — uniform `shuffle=True` as before. With them, the dataset prints a one-line summary at startup so you can verify the math:
+
+```
+[MRReportDataset] Rebalancing enabled (strategy=inverse_freq): 97896/97896 subjects matched labels CSV, weight range=[1.0, 1418.9], mean=33.0
+[Trainer] Using WeightedRandomSampler (strategy=inverse_freq)
+```
+
+Total compute is unchanged (still `--num_train_steps` steps × batch_size × num_processes draws). Rebalancing only **redistributes** that draw budget toward rare-pathology subjects via `WeightedRandomSampler(..., replacement=True)`.
 
 ### Training Configuration
 

contrastive-pretraining/scripts/data.py

@@ -6,10 +6,13 @@
 import torch
 import torch.nn.functional as F
 import nibabel as nib
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, WeightedRandomSampler
 from tqdm import tqdm
 
 
+REBALANCE_STRATEGIES = ('inverse_freq', 'sqrt_inverse_freq', 'max_inverse_freq')
+
+
 def cycle(dl):
     """Helper to infinitely loop through a DataLoader."""
     while True:
@@ -91,6 +94,16 @@ def normalize(self, data):
 }
 
 
+# Mapping from logical space name to the image subdirectory used in the
+# raw HuggingFace download layout (layout 2). The native repo stores volumes
+# in `img/`, while derivative repos (coreg, atlas) use prefixed names.
+SPACE_TO_IMG_SUBDIR = {
+    'native_space': 'img',
+    'coreg_space': 'coreg_img',
+    'atlas_space': 'atlas_img',
+}
+
+
 class MRReportDataset(Dataset):
     """
     Dataset for brain MRI with variable numbers of volumes per subject.
@@ -120,6 +133,10 @@ def __init__(
         normalizer_kwargs=None,
         splits_csv=None,
         split="train",
+        pathology_labels_csv=None,
+        rebalance_strategy=None,
+        rebalance_base_weight=1.0,
+        rebalance_eps=1e-6,
     ):
         self.data_folder = data_folder
         self.space = space
@@ -144,6 +161,17 @@ def __init__(
         # Discover subjects
         self.samples = self._prepare_samples(data_folder)
 
+        # Optional inverse-prevalence rebalancing weights for rare pathologies
+        self.rebalance_strategy = rebalance_strategy
+        self.label_columns = []
+        self.label_prevalence = None
+        self.sample_weights = self._compute_sample_weights(
+            pathology_labels_csv,
+            rebalance_strategy,
+            rebalance_base_weight,
+            rebalance_eps,
+        )
+
         print(f"[MRReportDataset] Found {len(self.samples)} subjects")
         for s in self.samples[:5]:
             print(f"  - {s['subject_id']}: {len(s['image_paths'])} volumes, {len(s['sentences'])} sentences")
@@ -182,7 +210,13 @@ def _prepare_samples(self, data_folder):
 
         Supports two directory layouts:
           1) data_folder/<study_uid>/<space>/img/*.nii.gz
-          2) data_folder/batchXX/<study_uid>/img/*.nii.gz  (legacy HuggingFace layout)
+          2) data_folder/batchXX/<study_uid>/<img_subdir>/*.nii.gz
+             (raw HuggingFace download layout; <img_subdir> depends on space)
+
+        For layout 2 the image subdirectory depends on the chosen space:
+          - native_space -> img/        (Forithmus/MR-RATE)
+          - coreg_space  -> coreg_img/  (Forithmus/MR-RATE-coreg)
+          - atlas_space  -> atlas_img/  (Forithmus/MR-RATE-atlas)
 
         Layout is auto-detected: if the first subdirectory contains a <space>
         subfolder, layout 1 is used; otherwise layout 2.
@@ -207,11 +241,12 @@ def _prepare_samples(self, data_folder):
                 img_dir = os.path.join(data_folder, study_uid, self.space, 'img')
                 self._add_subject(samples, study_uid, img_dir)
         else:
-            # Layout 2: data_folder/batchXX/<study_uid>/img/
+            # Layout 2: data_folder/batchXX/<study_uid>/<img_subdir>/
+            img_subdir = SPACE_TO_IMG_SUBDIR.get(self.space, 'img')
             for batch_dir in first_level_dirs:
                 batch_path = os.path.join(data_folder, batch_dir)
                 for study_uid in sorted(os.listdir(batch_path)):
-                    img_dir = os.path.join(batch_path, study_uid, 'img')
+                    img_dir = os.path.join(batch_path, study_uid, img_subdir)
                     self._add_subject(samples, study_uid, img_dir)
 
         return samples
@@ -238,6 +273,119 @@ def _add_subject(self, samples, study_uid, img_dir):
             'sentences': self.subject_to_sentences[study_uid],
         })
 
+    def _compute_sample_weights(self, csv_path, strategy, base_weight, eps):
+        """Compute per-subject sampling weights from a pathology-labels CSV.
+
+        Inverse-prevalence weighting upsamples subjects with rare positive
+        pathologies so contrastive batches see them more often. Subjects not
+        listed in the CSV (or all-negative subjects) receive `base_weight`.
+
+        Strategies:
+          - 'inverse_freq':       base + sum_p y_p * (1 / prevalence_p)
+          - 'sqrt_inverse_freq':  base + sum_p y_p * sqrt(1 / prevalence_p)
+          - 'max_inverse_freq':   max(base, max_p y_p * (1 / prevalence_p))
+
+        Returns:
+            A torch.FloatTensor of length len(self.samples) if rebalancing is
+            enabled, else None. Weights are unnormalized (WeightedRandomSampler
+            normalizes internally).
+        """
+        if csv_path is None or strategy is None:
+            return None
+        if strategy not in REBALANCE_STRATEGIES:
+            raise ValueError(
+                f"Unknown rebalance_strategy '{strategy}'. "
+                f"Choose from: {list(REBALANCE_STRATEGIES)}"
+            )
+
+        labels_by_uid, label_columns = self._load_pathology_labels(csv_path)
+        self.label_columns = label_columns
+
+        # Compute prevalence over the subset of dataset subjects that have labels
+        label_rows = [labels_by_uid[s['subject_id']] for s in self.samples
+                      if s['subject_id'] in labels_by_uid]
+        if not label_rows:
+            print(
+                f"[MRReportDataset] WARNING: no dataset subjects matched the "
+                f"pathology labels CSV; rebalancing disabled."
+            )
+            return None
+        label_matrix = np.stack(label_rows, axis=0)
+        prevalence = label_matrix.mean(axis=0)
+        self.label_prevalence = prevalence
+        inv_freq = 1.0 / np.clip(prevalence, eps, None)
+
+        if strategy == 'sqrt_inverse_freq':
+            per_class = np.sqrt(inv_freq)
+        else:
+            per_class = inv_freq
+
+        weights = np.full(len(self.samples), base_weight, dtype=np.float32)
+        for i, s in enumerate(self.samples):
+            y = labels_by_uid.get(s['subject_id'])
+            if y is None:
+                continue
+            if strategy == 'max_inverse_freq':
+                pos = y * inv_freq
+                weights[i] = max(base_weight, float(pos.max()))
+            else:
+                weights[i] = base_weight + float((y * per_class).sum())
+
+        n_labeled = sum(1 for s in self.samples if s['subject_id'] in labels_by_uid)
+        print(
+            f"[MRReportDataset] Rebalancing enabled (strategy={strategy}): "
+            f"{n_labeled}/{len(self.samples)} subjects matched labels CSV, "
+            f"weight range=[{weights.min():.3g}, {weights.max():.3g}], "
+            f"mean={weights.mean():.3g}"
+        )
+        return torch.from_numpy(weights)
+
+    @staticmethod
+    def _load_pathology_labels(csv_path):
+        """Load a pathology labels CSV.
+
+        Expects a header row with 'study_uid' (or 'subject_id') followed by
+        one binary column per pathology. Returns (dict uid -> np.ndarray,
+        list of label column names).
+        """
+        labels = {}
+        with open(csv_path, 'r') as f:
+            reader = csv.DictReader(f)
+            fields = reader.fieldnames or []
+            if 'study_uid' in fields:
+                id_col = 'study_uid'
+            elif 'subject_id' in fields:
+                id_col = 'subject_id'
+            else:
+                raise ValueError(
+                    f"Pathology labels CSV {csv_path} must have a 'study_uid' "
+                    f"or 'subject_id' column. Got: {fields}"
+                )
+            label_columns = [c for c in fields if c != id_col]
+            for row in reader:
+                labels[row[id_col]] = np.array(
+                    [float(row[c]) for c in label_columns], dtype=np.float32
+                )
+        return labels, label_columns
+
+    def get_weighted_sampler(self, num_samples=None, generator=None):
+        """Build a WeightedRandomSampler from the precomputed sample weights.
+
+        Replacement sampling is required so high-weight (rare-pathology)
+        subjects can be drawn multiple times per epoch.
+        """
+        if self.sample_weights is None:
+            raise RuntimeError(
+                "sample_weights not computed; pass pathology_labels_csv and "
+                "rebalance_strategy to the dataset constructor."
+            )
+        return WeightedRandomSampler(
+            weights=self.sample_weights,
+            num_samples=num_samples or len(self.samples),
+            replacement=True,
+            generator=generator,
+        )
+
     def __len__(self):
         return len(self.samples)
 

contrastive-pretraining/scripts/data_inference.py

@@ -29,7 +29,7 @@
 import nibabel as nib
 from torch.utils.data import Dataset
 
-from data import NORMALIZERS, resize_array
+from data import NORMALIZERS, SPACE_TO_IMG_SUBDIR, resize_array
 
 
 class MRReportDatasetInfer(Dataset):
@@ -128,7 +128,8 @@ def _prepare_samples(self, data_folder):
 
         Supports two directory layouts (same auto-detection as training):
           1) data_folder/<study_uid>/<space>/img/*.nii.gz
-          2) data_folder/batchXX/<study_uid>/img/*.nii.gz
+          2) data_folder/batchXX/<study_uid>/<img_subdir>/*.nii.gz
+             (img_subdir is space-dependent: img/, coreg_img/, atlas_img/)
         """
         samples = []
 
@@ -147,10 +148,11 @@ def _prepare_samples(self, data_folder):
                 img_dir = os.path.join(data_folder, study_uid, self.space, 'img')
                 self._add_subject(samples, study_uid, img_dir)
         else:
+            img_subdir = SPACE_TO_IMG_SUBDIR.get(self.space, 'img')
             for batch_dir in first_level_dirs:
                 batch_path = os.path.join(data_folder, batch_dir)
                 for study_uid in sorted(os.listdir(batch_path)):
-                    img_dir = os.path.join(batch_path, study_uid, 'img')
+                    img_dir = os.path.join(batch_path, study_uid, img_subdir)
                     self._add_subject(samples, study_uid, img_dir)
 
         return samples