Add MR DINO online/offline MIL and vision-language training

README.md

@@ -40,6 +40,10 @@ Each folder includes its own `README.md` detailing configuration, dependencies,
 
 ## Workflow Summary
 
+For the Dong MR DINO checkpoints, full multi-sequence MIL (online/offline) and
+streaming contrastive VLM are documented in
+[`dino-pretraining/MIL_VLM.md`](dino-pretraining/MIL_VLM.md).
+
 1. **Download preprocessed data** directly from [MR-RATE on Hugging Face](https://huggingface.co/datasets/Forithmus/MR-RATE) (see `data-preprocessing/` for preprocessing details)
 2. **Train the contrastive model** on (multi-sequence MRI, radiology report) pairs using `contrastive-pretraining/`
 4. **Run zero-shot inference** for brain MRI pathology classification using trained model weights and `contrastive-pretraining/`

contrastive-pretraining/README.md

@@ -7,7 +7,7 @@ Contrastive vision-language model for brain & spine MRI that aligns multi-sequen
 
 ## Architecture
 
-Built on [FORA](https://github.com/forithmus/FORA), extending its image encoder and contrastive pretraining framework to multi-sequence brain & spine MRI with variable volumes per subject.
+Built on [MR-RATE](https://github.com/forithmus/MR-RATE), extending its image encoder and contrastive pretraining framework to multi-sequence brain & spine MRI with variable volumes per subject.
 
 **Text Encoder**: [BiomedVLP-CXR-BERT-specialized](https://huggingface.co/microsoft/BiomedVLP-CXR-BERT-specialized).
 

contrastive-pretraining/scripts/data.py

@@ -109,14 +109,21 @@ def normalize(self, data):
 CACHE_MANIFEST_NAME = "_manifest.json"
 
 
+def is_mri_image_file(filename):
+    name = str(filename).lower()
+    auxiliary = ("_brain-mask.nii.gz", "_defacing-mask.nii.gz", "_mask.nii.gz",
+                 "_seg.nii.gz", "_segmentation.nii.gz")
+    return name.endswith(".nii.gz") and not name.endswith(auxiliary)
+
+
 def list_nii_files(img_dir):
     """Sorted list of absolute *.nii.gz paths in img_dir (empty if it's missing)."""
     if not os.path.isdir(img_dir):
         return []
     return sorted(
         os.path.join(img_dir, f)
         for f in os.listdir(img_dir)
-        if f.endswith('.nii.gz')
+        if is_mri_image_file(f)
     )
 
 

contrastive-pretraining/scripts/extract_features.py

@@ -21,7 +21,7 @@
 seconds — no need to re-encode the 3D volumes every epoch.
 
 Token mode stops immediately before MR-RATE's final masked mean. It is the
-required input for NeuroVFM-style Classify-Then-Aggregate MIL; pooled
+required input for class-specific Classify-Then-Aggregate MIL; pooled
 `features_*.npy` cannot be used for MIL because the instances are gone.
 
 Usage:

contrastive-pretraining/scripts/mil_probe.py

@@ -1,7 +1,7 @@
 """
-NeuroVFM Classify-Then-Aggregate MIL on frozen MR-RATE token features.
+MRI Classify-Then-Aggregate MIL on frozen MR-RATE token features.
 
-The head matches NeuroVFM's diagnostic architecture while replacing only
+The head matches the diagnostic architecture while replacing only
 its visual encoder output with projected tokens from a contrastively
 pretrained MR-RATE checkpoint:
 
@@ -11,7 +11,7 @@
 
 The encoder is frozen by construction: extract_features.py runs it once
 under no_grad/eval and this script trains only the MIL head. Standard
-nn.Linear is used instead of NeuroVFM's FlashAttention FusedDense; the
+nn.Linear is used instead of the FlashAttention FusedDense; the
 parameterization and computation are otherwise the same.
 
 Workflow:
@@ -45,7 +45,7 @@
 from torch.utils.data import DataLoader, Dataset
 
 class PatchClassifier(nn.Module):
-    """NeuroVFM's shared MLP: Linear/GELU hidden blocks, then Linear."""
+    """the shared MLP: Linear/GELU hidden blocks, then Linear."""
 
     def __init__(self, dim: int, hidden_dims: tuple[int, ...], out_dim: int):
         super().__init__()
@@ -62,11 +62,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class ClassifyThenAggregate(nn.Module):
-    """NeuroVFM's class-specific gated-attention MIL head.
+    """the class-specific gated-attention MIL head.
 
     Inputs are flat tokens from a batch of ragged bags plus cumulative bag
     boundaries. Softmax is independently normalized within each bag and
-    class, exactly matching NeuroVFM's segment_csr implementation.
+    class, exactly matching the segment_csr implementation.
     """
 
     def __init__(
@@ -439,7 +439,7 @@ def factor(step: int) -> float:
 
 def main() -> None:
     parser = argparse.ArgumentParser(
-        "NeuroVFM Classify-Then-Aggregate MIL on frozen MR-RATE tokens"
+        "MRI Classify-Then-Aggregate MIL on frozen MR-RATE tokens"
     )
     parser.add_argument("--features_dir", type=str, required=True)
     parser.add_argument("--results_dir", type=str, default="./mil_probe_results")
@@ -543,7 +543,7 @@ def main() -> None:
         use_output_bias_scale=True,
     ).to(device)
     print(
-        f"NeuroVFM CTA head: dim={train_data.dim}, attention_hidden={args.hidden_dim}, "
+        f"MRI CTA head: dim={train_data.dim}, attention_hidden={args.hidden_dim}, "
         f"patch_mlp=[{args.mlp_hidden_dim}], classes={n_classes}"
     )
     print(
@@ -635,7 +635,7 @@ def main() -> None:
     print(f"Test mean AUROC: {test_mean:.4f}")
 
     architecture = {
-        "name": "NeuroVFM ClassifyThenAggregate",
+        "name": "MRI ClassifyThenAggregate",
         "dim": train_data.dim,
         "n_classes": n_classes,
         "hidden_dim": args.hidden_dim,

data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_neurovfm_diagnosis_extraction/README.md → data-preprocessing/src/mr_rate_preprocessing/reports_preprocessing/07_mri_diagnosis_extraction/README.md

@@ -1,8 +1,8 @@
-# NeuroVFM MRI diagnosis extraction
+# MRI MRI diagnosis extraction
 
 This is the report-to-label pipeline used for the current MR-RATE clinical
 evaluation target. It classifies each report against the 74 expert-defined
-NeuroVFM MRI diagnoses and returns a binary label plus a short rationale for
+MRI MRI diagnoses and returns a binary label plus a short rationale for
 every diagnosis.
 
 The checked-in prompt, diagnosis order, guidance, JSON parser, and inference
@@ -13,15 +13,15 @@ the repository only to reproduce the historical 37-pathology labels.
 Run one or more Slurm ranks with:
 
 ```bash
-srun --gpu-bind=none python extract_neurovfm_dx_gemma.py \
+srun --gpu-bind=none python extract_mri_dx_gemma.py \
   --reports_dir /path/to/reports \
-  --diagnoses_json data/neurovfm_mri_diagnoses.json \
+  --diagnoses_json data/mri_diagnoses_74.json \
   --output_dir /path/to/shards
 
 python merge_labels.py \
   --input_dir /path/to/shards \
-  --output /path/to/neurovfm74_labels.csv \
-  --rationales /path/to/neurovfm74_rationales.json
+  --output /path/to/mri74_labels.csv \
+  --rationales /path/to/mri74_rationales.json
 ```
 
 Input files are named `batchNN_reports.csv` and require `study_uid`. The