Fix/fora v14 recipe

dino-pretraining/mr_dino/__init__.py

@@ -1,5 +1,5 @@
-"""Volumetric DINOv3 training for atlas-registered MR-RATE studies."""
+"""Volumetric DINOv3 training for aligned (coreg/atlas) MR-RATE studies."""
 
-from .data import CropSpec, MRAtlasDINO3DDataset, collate_dino3d
+from .data import CropSpec, MRAlignedDINO3DDataset, collate_dino3d
 
-__all__ = ["CropSpec", "MRAtlasDINO3DDataset", "collate_dino3d"]
+__all__ = ["CropSpec", "MRAlignedDINO3DDataset", "collate_dino3d"]

dino-pretraining/mr_dino/data.py

@@ -1,6 +1,10 @@
-"""Atlas-registered multi-sequence MR sampling for volumetric DINOv3.
+"""Aligned multi-sequence MR sampling for volumetric DINOv3.
 
-The default path discovers and preprocesses atlas NIfTIs with the canonical
+Training uses MR-RATE's co-registered space by default (``coreg_space``: every
+sequence rigidly registered to the study's T1w center scan and kept on that
+scan's native grid, so no MNI field-of-view crop and no 1 mm resampling);
+``atlas_space`` remains selectable.  Both are voxel-aligned across sequences.
+The default path discovers and preprocesses NIfTIs with the canonical
 loader used by the previous MR-RATE training code.  The optional NPZ path reads
 the volume cache produced by that loader's ``preprocess_volumes.py`` helper.
 Global and local DINO views share a physical region, but may use different
@@ -25,8 +29,12 @@
 import torch.nn.functional as F
 from torch.utils.data import Dataset, Sampler
 
+from .recipe import overlap_pairs, snap_second_start, token_offset
+
 
 CACHE_MANIFEST_NAME = "_manifest.json"
+ALIGNED_SPACES = ("coreg_space", "atlas_space")
+DEFAULT_SPACE = "coreg_space"
 EXPECTED_CACHE_LAYOUT = "per_subject_stack"
 _PREVIOUS_DATA_MODULE = None
 
@@ -209,21 +217,30 @@ def _cache_space_dir(preprocessed_dir: str, space: str) -> Path:
     return nested if nested.is_dir() else root
 
 
-def validate_atlas_cache(
+def check_aligned_space(space: str) -> str:
+    """Cross-sequence views need voxel alignment: native_space is never accepted."""
+    if space not in ALIGNED_SPACES:
+        raise ValueError(
+            f"MR DINO requires an aligned space {ALIGNED_SPACES}; got {space!r} "
+            "(native_space sequences are not voxel-aligned)"
+        )
+    return space
+
+
+def validate_aligned_cache(
     preprocessed_dir: str,
-    space: str = "atlas_space",
+    space: str = DEFAULT_SPACE,
     expected_spacing: tuple[float, float, float] = (1.0, 0.5, 0.5),
 ) -> dict:
-    """Fail early unless this is a compatible MR-RATE atlas-space cache."""
-    if space != "atlas_space":
-        raise ValueError("MR DINO requires space='atlas_space'; native/coreg inputs are not accepted")
+    """Fail early unless this is a compatible MR-RATE aligned-space cache."""
+    check_aligned_space(space)
     space_dir = _cache_space_dir(preprocessed_dir, space)
     manifest_path = space_dir / CACHE_MANIFEST_NAME
     if not manifest_path.is_file():
         raise FileNotFoundError(f"Missing MR-RATE preprocessing manifest: {manifest_path}")
     manifest = json.loads(manifest_path.read_text())
-    if manifest.get("space") != "atlas_space":
-        raise ValueError(f"Cache is {manifest.get('space')!r}, expected 'atlas_space'")
+    if manifest.get("space") != space:
+        raise ValueError(f"Cache is {manifest.get('space')!r}, expected {space!r}")
     if manifest.get("layout") != EXPECTED_CACHE_LAYOUT:
         raise ValueError(
             f"Unsupported cache layout {manifest.get('layout')!r}; expected {EXPECTED_CACHE_LAYOUT!r}"
@@ -234,13 +251,13 @@ def validate_atlas_cache(
     return manifest
 
 
-def discover_atlas_cache(
+def discover_aligned_cache(
     preprocessed_dir: str,
     splits_csv: str | None = None,
     split: str = "train",
-    space: str = "atlas_space",
+    space: str = DEFAULT_SPACE,
 ) -> list[dict]:
-    manifest = validate_atlas_cache(preprocessed_dir, space)
+    manifest = validate_aligned_cache(preprocessed_dir, space)
     space_dir = _cache_space_dir(preprocessed_dir, space)
     selected = _load_split_ids(splits_csv, split)
     samples = []
@@ -256,7 +273,7 @@ def discover_atlas_cache(
             "volume_shape": shape[1:],
         })
     if not samples:
-        raise RuntimeError(f"No {split!r} atlas-registered studies found under {space_dir}")
+        raise RuntimeError(f"No {split!r} {space} studies found under {space_dir}")
     return samples
 
 
@@ -306,14 +323,16 @@ def previous_mr_data_module():
     return module
 
 
-def discover_raw_atlas(
+def discover_raw_aligned(
     data_folder: str,
     selected: set[str] | None = None,
+    space: str = DEFAULT_SPACE,
 ) -> list[dict]:
-    """Discover atlas NIfTIs through the previous training dataloader."""
+    """Discover aligned NIfTIs (coreg_img/ or atlas_img/) through the previous training dataloader."""
+    check_aligned_space(space)
     samples = []
     previous = previous_mr_data_module()
-    for subject in previous.discover_subjects(data_folder, "atlas_space"):
+    for subject in previous.discover_subjects(data_folder, space):
         study_uid = str(subject["subject_id"])
         if selected is not None and study_uid not in selected:
             continue
@@ -324,17 +343,18 @@ def discover_raw_atlas(
             "n_sequences": len(paths),
         })
     if not samples:
-        raise RuntimeError(f"No atlas-registered NIfTIs found under {data_folder}")
+        raise RuntimeError(f"No {space} NIfTIs found under {data_folder}")
     return samples
 
 
-def discover_raw_atlas_split(
+def discover_raw_aligned_split(
     data_folder: str,
     splits_csv: str | None = None,
     split: str = "train",
+    space: str = DEFAULT_SPACE,
 ) -> list[dict]:
-    """Discover a split using the previous loader's atlas-space convention."""
-    return discover_raw_atlas(data_folder, _load_split_ids(splits_csv, split))
+    """Discover a split using the previous loader's space convention."""
+    return discover_raw_aligned(data_folder, _load_split_ids(splits_csv, split), space)
 
 
 def _raw_volume(
@@ -356,8 +376,8 @@ def _raw_volume(
     return torch.from_numpy(np.ascontiguousarray(array)).to(torch.bfloat16)
 
 
-class MRAtlasDINO3DDataset(Dataset):
-    """Study-level SSL samples from atlas-aligned, variable-count MR sequences."""
+class MRAlignedDINO3DDataset(Dataset):
+    """Study-level SSL samples from voxel-aligned, variable-count MR sequences."""
 
     def __init__(
         self,
@@ -366,23 +386,31 @@ def __init__(
         data_folder: str | None = None,
         splits_csv: str | None = None,
         split: str = "train",
-        space: str = "atlas_space",
+        space: str = DEFAULT_SPACE,
         crop_spec: CropSpec = CropSpec(),
         target_spacing: tuple[float, float, float] = (1.0, 0.5, 0.5),
         target_shape: tuple[int, int, int] = (256, 384, 384),
         posterior_shift_mm: float = 15.0,
         cache_files: list[str] | None = None,
         raw_samples: list[dict] | None = None,
-        cross_sequence_probability: float = 0.75,
+        cross_sequence_probability: float = 0.25,
         candidate_trials: int = 12,
         seed: int = 3407,
+        global_overlap: float = 0.25,
+        patch_size: tuple[int, int, int] = (2, 16, 16),
     ) -> None:
         if (preprocessed_dir is None) == (data_folder is None):
             raise ValueError("Pass exactly one of preprocessed_dir or data_folder")
-        if space != "atlas_space":
-            raise ValueError("Only atlas_space is valid for atlas-registered MR DINO")
+        self.space = check_aligned_space(space)
         if not 0 <= cross_sequence_probability <= 1:
             raise ValueError("cross_sequence_probability must be in [0, 1]")
+        if not 0 <= global_overlap <= 1:
+            raise ValueError("global_overlap must be in [0, 1]")
+        # FORA fix: the two global crops only need >= 25% overlap per axis (0.75 made them near
+        # duplicates), and the second crop is snapped to the patch lattice so that every token in
+        # the overlap has an exact twin for the cross-crop patch objective.
+        self.global_overlap = float(global_overlap)
+        self.patch_size = tuple(int(v) for v in patch_size)
         self.crop_spec = crop_spec
         self.target_shape = tuple(int(x) for x in target_shape)
         self.target_spacing = tuple(float(x) for x in target_spacing)
@@ -395,10 +423,10 @@ def __init__(
         self._cached_names: list[str] | None = None
         self.preprocessed = preprocessed_dir is not None
         if self.preprocessed:
-            manifest = validate_atlas_cache(preprocessed_dir, space, target_spacing)
+            manifest = validate_aligned_cache(preprocessed_dir, space, target_spacing)
             self.target_shape = tuple(int(x) for x in manifest["target_shape"])
             if cache_files is None:
-                self.samples = discover_atlas_cache(preprocessed_dir, splits_csv, split, space)
+                self.samples = discover_aligned_cache(preprocessed_dir, splits_csv, split, space)
             else:
                 selected = _load_split_ids(splits_csv, split)
                 self.samples = [
@@ -412,18 +440,18 @@ def __init__(
                     if selected is None or Path(path).stem in selected
                 ]
                 if not self.samples:
-                    raise RuntimeError("This rank received no atlas-registered cache studies")
+                    raise RuntimeError(f"This rank received no {space} cache studies")
         else:
             selected = _load_split_ids(splits_csv, split)
             if raw_samples is None:
-                self.samples = discover_raw_atlas(data_folder, selected)
+                self.samples = discover_raw_aligned(data_folder, selected, space)
             else:
                 self.samples = [
                     sample for sample in raw_samples
                     if selected is None or sample["study_uid"] in selected
                 ]
                 if not self.samples:
-                    raise RuntimeError("This rank received no atlas-registered NIfTI studies")
+                    raise RuntimeError(f"This rank received no {space} NIfTI studies")
         for sample in self.samples:
             if "volume_shape" in sample and tuple(sample["volume_shape"]) != self.target_shape:
                 raise ValueError(
@@ -500,10 +528,13 @@ def __getitem__(self, key: int | tuple[int, int]) -> dict:
         first_start = _informative_start(
             stack[first_sequence], self.crop_spec.global_shape, rng, self.candidate_trials
         )
-        global_starts = _overlapping_global_starts(
-            tuple(stack.shape[1:]), self.crop_spec.global_shape, rng, first=first_start
+        container = tuple(stack.shape[1:])
+        first, second = _overlapping_global_starts(
+            container, self.crop_spec.global_shape, rng, first=first_start,
+            min_overlap=self.global_overlap,
         )
-        lo, hi = _intersection_box(global_starts, self.crop_spec.global_shape)
+        second = snap_second_start(first, second, container, self.crop_spec.global_shape, self.patch_size)
+        global_starts = (first, second)
 
         teacher_globals, student_globals = [], []
         for sequence, start in zip(global_sequences, global_starts):
@@ -514,6 +545,11 @@ def __getitem__(self, key: int | tuple[int, int]) -> dict:
         local_views, local_sequences, local_starts = [], [], []
         for _ in range(self.crop_spec.local_crops):
             sequence = rng.randrange(n_sequences)
+            # FORA fix: locals lie inside ONE of the globals (the intersection of two crops with
+            # 25% overlap can be thinner than a local crop).
+            parent = global_starts[rng.randrange(2)]
+            lo = parent
+            hi = tuple(min(p + g, c) for p, g, c in zip(parent, self.crop_spec.global_shape, container))
             start = _contained_start(lo, hi, self.crop_spec.local_shape, rng)
             clean = _crop_at_start(stack[sequence], self.crop_spec.local_shape, start)
             local_views.append(_student_distortion(clean, rng).unsqueeze(0))
@@ -573,6 +609,7 @@ def collate_dino3d(
     samples: list[dict],
     patch_size: tuple[int, int, int] = (2, 16, 16),
     mask_ratio: tuple[float, float] = (0.1, 0.5),
+    cross_view_sequences: str = "same",
 ) -> dict:
     teacher = torch.stack([sample["teacher_global"] for sample in samples], dim=1)
     student = torch.stack([sample["student_global"] for sample in samples], dim=1)
@@ -590,7 +627,27 @@ def collate_dino3d(
     masks = torch.stack(masks)
     indices = masks.flatten().nonzero().flatten()
     weights = (1.0 / masks.sum(-1).clamp(min=1)).unsqueeze(-1).expand_as(masks)[masks]
+    # Twin tokens of the two global crops (cross-crop patch objective). By default only when both
+    # globals show the SAME sequence, preserving this module's contract that cross-sequence images
+    # are never patch-level targets; "any" also pairs co-registered voxels across contrasts.
+    if cross_view_sequences not in {"same", "any", "none"}:
+        raise ValueError("cross_view_sequences must be 'same', 'any' or 'none'")
+    n = math.prod(grid)
+    pa, pb, ps = [], [], []
+    for index, sample in enumerate(samples):
+        seq0, seq1 = sample["global_sequences"]
+        if cross_view_sequences == "none" or (cross_view_sequences == "same" and seq0 != seq1):
+            continue
+        g0, g1 = sample["global_starts"]
+        a, b = overlap_pairs(token_offset(g0, g1, patch_size), grid)
+        pa.append(a + index * n)
+        pb.append(b + index * n)
+        ps.append(torch.full((a.numel(),), index, dtype=torch.long))
+    empty = torch.empty(0, dtype=torch.long)
     return {
+        "cross_a": torch.cat(pa) if pa else empty,
+        "cross_b": torch.cat(pb) if pb else empty,
+        "cross_sample": torch.cat(ps) if ps else empty,
         "teacher_global": teacher,
         "student_global": student,
         "student_local": local,