diff --git a/unit_tests/test_models.py b/unit_tests/test_models.py index 5941822..55df3c7 100644 --- a/unit_tests/test_models.py +++ b/unit_tests/test_models.py @@ -113,3 +113,49 @@ def test_patch_centres_are_rounded_not_truncated(): patches = extractor.extract_patches(volume, torch.tensor([[[24.6, 24.6, 24.6]]])) centre = patches.shape[-1] // 2 assert patches[0, 0, 0, centre, centre, centre] == 1.0 + + +@pytest.mark.parametrize("weight_features", [True, False]) +def test_heatmap_densenet_landmark_features_match_the_broadcast_formulation(weight_features): + """The einsum contraction must equal the (B, N, F, H, W, D) broadcast it replaced. + + The broadcast is the readable definition; the einsum is the one that fits in memory + once ``n_landmarks`` or the heatmap resolution grows. Only the memory profile may + differ, never the numbers. + """ + from verpex.models.densenet import HeatmapDenseNet + + torch.manual_seed(0) + model = HeatmapDenseNet( + spatial_dims=3, + in_channels=1, + n_landmarks=5, + feature_l=7, + init_features=8, + growth_rate=4, + block_config=(2, 2), + weight_features=weight_features, + ).eval() + + with torch.no_grad(): + heatmaps, landmark_features, feature_map = model(torch.randn(2, 1, 32, 32, 32)) + + if weight_features: + expected = (heatmaps.unsqueeze(2) * feature_map.unsqueeze(1)).sum(dim=(3, 4, 5)) + else: + expected = feature_map.mean(dim=(2, 3, 4)).unsqueeze(1).expand(-1, heatmaps.shape[1], -1) + assert landmark_features.shape == (2, 5, 7) + assert torch.allclose(landmark_features, expected, atol=1e-5) + + +def test_heatmap_densenet_landmark_features_do_not_backpropagate_into_the_heatmaps(): + """The weighting is detached, so the feature encoding must not carry heatmap grads.""" + from verpex.models.densenet import HeatmapDenseNet + + torch.manual_seed(0) + model = HeatmapDenseNet( + spatial_dims=3, in_channels=1, n_landmarks=3, feature_l=4, init_features=8, growth_rate=4, block_config=(2, 2) + ) + heatmaps, landmark_features, _ = model(torch.randn(1, 1, 32, 32, 32)) + grad = torch.autograd.grad(landmark_features.sum(), heatmaps, allow_unused=True, retain_graph=True)[0] + assert grad is None diff --git a/unit_tests/test_transforms.py b/unit_tests/test_transforms.py index d18a0fd..dcec235 100644 --- a/unit_tests/test_transforms.py +++ b/unit_tests/test_transforms.py @@ -10,7 +10,14 @@ import pytest import torch -from verpex.data.transforms import LandMarksRandHorizontalFlipNeighbor +from verpex.data.transforms import ( + LandMarksRandHorizontalFlipNeighbor, + LandmarksRandLabelDropout, + LandmarksRandLimitedFov, + RandLabelCollapse, + create_transforms, +) +from verpex.registry import UnknownTypeError #: 81/82 are a left/right pair; the rest are midline landmarks that map to themselves. FLIP_PAIRS = {81: 82, 82: 81, 90: 90, 91: 91} @@ -100,3 +107,270 @@ def test_flipping_twice_restores_the_original(): out = flip(flip(batch)) assert torch.equal(out["target"], before_target) assert torch.equal(out["input"], before_input) + + +# --------------------------------------------------------------------------- +# Limited field of view, label dropout, label collapse, and the registry. +# --------------------------------------------------------------------------- + +FOV_SHAPE = (16, 20, 24) + + +def make_fov_sample(shape=FOV_SHAPE, label=3.0): + """Build a sample whose labelled bar hugs the high face of axis 2. + + Channel 0 is the label map, channel 1 a surface-like copy. The landmarks run along + the bar, one per slice, so it is obvious which of them a cut removes. + """ + volume = torch.zeros(2, *shape) + volume[0, 4:12, 6:14, 10 : shape[2]] = label + volume[1] = (volume[0] > 0).float() + target = torch.tensor([[8.0, 10.0, float(z)] for z in range(10, shape[2])]) + return { + "input": volume, + "target": target, + "target_indices": torch.arange(len(target)), + "loss_mask": torch.ones(len(target), dtype=torch.bool), + "label_channel": 0, + "landmark_labels": torch.full((len(target),), int(label)), + "sample_index": 0, + } + + +def clone_sample(sample, **overrides): + """Deep-enough copy so a transform's in-place writes cannot leak between cases.""" + out = {k: (v.clone() if torch.is_tensor(v) else v) for k, v in sample.items()} + out.update(overrides) + return out + + +def cut_axis2_at_both_ends(**kwargs): + """A fully determined transform: it removes 20% of each end of axis 2 and nothing else. + + Both ends rather than one because which single end gets cut is itself a random draw, + and these tests want the cut fixed, not the seed. + """ + params = { + "prob": (0.0, 0.0, 1.0), + "cut_range": ((0.0, 0.0), (0.0, 0.0), (0.2, 0.2)), + "both_sides_prob": 1.0, + "min_visible_landmarks": 1, + } + params.update(kwargs) + return LandmarksRandLimitedFov(**params) + + +def test_limited_fov_keeps_the_spatial_shape(): + """The model's input size is fixed, so the transform may zero but never resize.""" + sample = make_fov_sample() + out = cut_axis2_at_both_ends()(clone_sample(sample)) + assert out["input"].shape == sample["input"].shape + + +def test_limited_fov_zeroes_the_slab_it_removes(): + """After the cut and the shift, no labelled voxel may survive from the removed slab.""" + sample = make_fov_sample() + out = cut_axis2_at_both_ends()(clone_sample(sample)) + kept = int((out["input"][0] > 0).sum()) + # 20% of 24 is 4.8 -> 5 slices removed from a 14-slice bar. + assert kept == int((sample["input"][0] > 0).sum()) * 9 // 14 + + +def test_limited_fov_does_not_wrap_content_around_the_volume_edge(): + """A circular shift would put anatomy back at the face it was pushed off. + + ``torch.roll`` is the obvious way to translate a volume and it is wrong here: the + landmark coordinate lands back *inside* the volume, on the wrong structure, so no + bounds check catches it. The re-centring shift is driven by the label channel, and it + moves the uncut axes too - which is where content actually can wrap. + + Here the label sits high on axis 0, so re-centring shifts by -4 along it, and a marker + fills the four slices that a wrap would deposit at the opposite face. + """ + volume = torch.zeros(2, *FOV_SHAPE) + volume[0, 10:15, 8:12, 8:16] = 3.0 + volume[1, 0:4] = 7.0 + sample = { + "input": volume, + "target": torch.tensor([[12.0, 10.0, 12.0]]), + "loss_mask": torch.ones(1, dtype=torch.bool), + "label_channel": 0, + } + out = cut_axis2_at_both_ends()(sample) + assert float(out["input"][1][-4:].sum()) == 0.0 + + +def test_limited_fov_moves_the_volume_and_the_landmarks_by_the_same_delta(): + """A sign error in the shift would leave the landmarks off their own anatomy.""" + out = cut_axis2_at_both_ends()(clone_sample(make_fov_sample())) + kept = out["target"][out["loss_mask"]].round().long() + on_label = out["input"][0][kept[:, 0], kept[:, 1], kept[:, 2]] + assert bool((on_label > 0).all()) + + +def test_limited_fov_centres_the_surviving_label_on_the_volume_centre(): + """Re-centring must reproduce what the offline cutout does to a truncated scan.""" + out = cut_axis2_at_both_ends()(clone_sample(make_fov_sample())) + indices = torch.nonzero(out["input"][0] > 0) + centre = ((indices.amin(0) + indices.amax(0)) // 2).tolist() + assert centre == [size // 2 for size in FOV_SHAPE] + + +def test_limited_fov_masks_landmarks_that_fall_outside_the_kept_slab(): + """Supervising a landmark whose anatomy was cut away teaches hallucination.""" + out = cut_axis2_at_both_ends()(clone_sample(make_fov_sample())) + assert int(out["loss_mask"].sum()) == 9 + + +def test_limited_fov_intersects_an_existing_loss_mask_rather_than_replacing_it(): + """A landmark the dataset already excluded must not come back supervised.""" + sample = make_fov_sample() + sample["loss_mask"][0] = False + out = cut_axis2_at_both_ends()(clone_sample(sample)) + assert not bool(out["loss_mask"][0]) + + +def test_limited_fov_writes_a_loss_mask_when_the_sample_has_none(): + """Verpex's own datasets build the mask after the transform, so it has to be carried.""" + sample = make_fov_sample() + del sample["loss_mask"] + out = cut_axis2_at_both_ends()(sample) + assert out["loss_mask"].dtype == torch.bool + assert int(out["loss_mask"].sum()) == 9 + + +def test_limited_fov_is_a_noop_when_no_axis_is_cut(): + """Re-centring unconditionally would undo the affine's deliberate off-centre shift.""" + sample = make_fov_sample() + out = LandmarksRandLimitedFov(prob=(0.0, 0.0, 0.0))(clone_sample(sample)) + assert torch.equal(out["input"], sample["input"]) + assert torch.equal(out["target"], sample["target"]) + + +def test_limited_fov_is_a_noop_when_too_few_landmarks_would_survive(): + """A sample whose supervision the cut would wipe out is worth nothing; leave it whole.""" + sample = make_fov_sample() + out = cut_axis2_at_both_ends(min_visible_landmarks=999)(clone_sample(sample)) + assert torch.equal(out["input"], sample["input"]) + assert torch.equal(out["loss_mask"], sample["loss_mask"]) + + +def test_limited_fov_reads_the_label_channel_the_sample_names(): + """The stack order depends on input_data_type, so a hardcoded -1 is silently wrong. + + Here channel 1 holds a bar offset from channel 0's. Re-centring on the named channel + must put *that* bar on the volume centre. + """ + sample = make_fov_sample() + sample["input"][1] = 0.0 + sample["input"][1, 2:6, 2:6, 0:8] = 5.0 + out = cut_axis2_at_both_ends(label_channel=0)(clone_sample(sample, label_channel=1)) + indices = torch.nonzero(out["input"][1] > 0) + assert ((indices.amin(0) + indices.amax(0)) // 2).tolist() == [size // 2 for size in FOV_SHAPE] + + +def test_limited_fov_with_a_seed_gives_the_same_cut_for_the_same_sample(): + """An evaluation metric is only comparable across epochs if the cut is repeatable.""" + transform = LandmarksRandLimitedFov(seed=11, min_visible_landmarks=1) + torch.manual_seed(1) + first = transform(clone_sample(make_fov_sample())) + torch.manual_seed(999) + second = transform(clone_sample(make_fov_sample())) + assert torch.equal(first["input"], second["input"]) + assert torch.equal(first["loss_mask"], second["loss_mask"]) + + +def test_limited_fov_with_a_seed_gives_a_different_cut_for_a_different_sample(): + """Every scan getting the identical cut would measure one truncation, not a range.""" + transform = LandmarksRandLimitedFov(seed=11, min_visible_landmarks=1) + outs = [transform(clone_sample(make_fov_sample(), sample_index=i)) for i in range(8)] + assert any(not torch.equal(outs[0]["input"], other["input"]) for other in outs[1:]) + + +def test_limited_fov_without_a_seed_varies_across_calls(): + """Training wants a fresh cut every epoch, which is what the global RNG gives.""" + transform = LandmarksRandLimitedFov(min_visible_landmarks=1) + torch.manual_seed(0) + outs = [transform(clone_sample(make_fov_sample())) for _ in range(8)] + assert any(not torch.equal(outs[0]["input"], other["input"]) for other in outs[1:]) + + +@pytest.mark.parametrize("bad", [{"prob": (0.5, 0.5)}, {"prob": (2.0, 0, 0)}, {"cut_range": ((0.6, 0.7),) * 3}, {"min_keep_fraction": 0}]) +def test_limited_fov_rejects_a_malformed_range(bad): + """A bad range would otherwise degrade to a silent no-op for a whole training run.""" + with pytest.raises(ValueError): + LandmarksRandLimitedFov(**bad) + + +def test_label_dropout_removes_the_label_from_every_channel(): + """A field of view that lost a rib lost its surface and its intensities too.""" + sample = make_fov_sample() + sample["input"][0, 0:3, 0:3, 0:3] = 9.0 + sample["input"][1, 0:3, 0:3, 0:3] = 1.0 + out = LandmarksRandLabelDropout(prob=1.0, n_labels=(1, 1), end_bias=1.0)(clone_sample(sample)) + present = torch.unique(out["input"][0]) + dropped = {3.0, 9.0} - set(present.tolist()) + assert len(dropped) == 1 + value = dropped.pop() + assert float(out["input"][1][sample["input"][0] == value].sum()) == 0.0 + + +def test_label_dropout_unsupervises_the_landmarks_of_the_label_it_removed(): + """Asking the model to find anatomy absent from its own input teaches hallucination.""" + sample = make_fov_sample() + sample["input"][0, 0:3, 0:3, 0:3] = 9.0 + out = LandmarksRandLabelDropout(prob=1.0, n_labels=(1, 1), end_bias=1.0)(clone_sample(sample)) + gone = 3.0 not in torch.unique(out["input"][0]).tolist() + assert int(out["loss_mask"].sum()) == (0 if gone else len(sample["target"])) + + +def test_label_dropout_never_empties_the_input(): + """A volume with no anatomy left has nothing to learn from.""" + sample = make_fov_sample() + out = LandmarksRandLabelDropout(prob=1.0, n_labels=(1, 5))(clone_sample(sample)) + assert bool((out["input"][0] > 0).any()) + + +def test_label_collapse_makes_every_label_value_identical(): + """The point is to remove the identity the model would otherwise read off the value.""" + sample = make_fov_sample() + sample["input"][0, 0:3, 0:3, 0:3] = 9.0 + out = RandLabelCollapse(prob=1.0, channels=[0])({"input": sample["input"].clone()}) + assert sorted(torch.unique(out["input"][0]).tolist()) == [0.0, 1.0] + + +def test_label_collapse_leaves_the_channels_it_was_not_given(): + """It must not reach a bookkeeping copy of the label map the loss mask comes from.""" + sample = make_fov_sample() + out = RandLabelCollapse(prob=1.0, channels=[0])({"input": sample["input"].clone()}) + assert torch.equal(out["input"][1], sample["input"][1]) + + +def test_label_collapse_without_channels_is_an_error_not_a_silent_noop(): + """Quietly doing nothing for a whole run is the worst outcome for an ablation.""" + with pytest.raises(KeyError, match="label channels"): + RandLabelCollapse(prob=1.0)({"input": torch.zeros(1, 4, 4, 4)}) + + +def test_a_single_transform_config_still_builds_one_transform(): + """Every config written before the pipeline became a list must keep working.""" + built = create_transforms({"type": "LandmarksRandLimitedFov"}) + assert len(built) == 1 + assert isinstance(built[0], LandmarksRandLimitedFov) + + +def test_a_list_of_transform_configs_builds_them_in_order(): + """Order is part of the meaning: the field-of-view cut has to run after the affine.""" + built = create_transforms([{"type": "LandmarksRandLimitedFov"}, {"type": "RandLabelCollapse"}]) + assert [type(t) for t in built] == [LandmarksRandLimitedFov, RandLabelCollapse] + + +def test_no_transform_config_builds_nothing(): + """The datasets always take a `transforms=` argument, so None must survive.""" + assert create_transforms(None) is None + + +def test_an_unregistered_transform_type_names_the_registered_ones(): + """The old factory raised a bare ValueError and could reach only two of the classes.""" + with pytest.raises(UnknownTypeError, match="LandmarksRandLimitedFov"): + create_transforms({"type": "NotATransform"}) diff --git a/verpex/data/dataset.py b/verpex/data/dataset.py index 9beded7..5c96da8 100644 --- a/verpex/data/dataset.py +++ b/verpex/data/dataset.py @@ -293,7 +293,10 @@ def __getitem__(self, index): missing_poi_list_idx = [self.poi_idx_to_list_idx[missing_poi.item()] for missing_poi in missing_pois] loss_mask[missing_poi_list_idx] = 0 - data_dict["loss_mask"] = loss_mask.bool() + # A transform may have masked landmarks itself - a limited-FOV cut removes the + # anatomy a landmark sits on. Intersect rather than overwrite, or that is lost. + carried = data_dict.get("loss_mask") + data_dict["loss_mask"] = loss_mask.bool() if carried is None else (loss_mask.bool() & carried.bool()) transformed_mask = mask if self.show_neighbors: @@ -792,7 +795,10 @@ def __getitem__(self, index): combined_loss_mask[outside_poi_indices] = 0 - data_dict["loss_mask"] = combined_loss_mask.bool() + # A transform may have masked landmarks itself - a limited-FOV cut removes the + # anatomy a landmark sits on. Intersect rather than overwrite, or that is lost. + carried = data_dict.get("loss_mask") + data_dict["loss_mask"] = combined_loss_mask.bool() if carried is None else (combined_loss_mask.bool() & carried.bool()) transformed_mask = torch.from_numpy(mask) surface = compute_surface(transformed_mask, iterations=self.iterations) diff --git a/verpex/data/transforms.py b/verpex/data/transforms.py index 5f4ce25..a6db3b8 100644 --- a/verpex/data/transforms.py +++ b/verpex/data/transforms.py @@ -56,6 +56,8 @@ from monai.utils.enums import TraceKeys, TransformBackends from monai.utils.type_conversion import convert_data_type, get_equivalent_dtype +from verpex.registry import build + class LandmarksRandAffine: def __init__( @@ -205,6 +207,421 @@ def __call__(self, dd): return dd +def _sample_generator(seed, dd, index_key): + """Return a per-sample RNG, or ``None`` to draw from the global one. + + Seeding from the sample's own index is what makes an evaluation-time augmentation + repeatable: ``worker_init_fn`` reseeds workers from ``torch.initial_seed()``, which + Lightning advances every epoch, so the global RNG cannot give a sample the same draw + twice. + """ + if seed is None: + return None + index = dd.get(index_key, 0) + if torch.is_tensor(index): + index = int(index.item()) + generator = torch.Generator() + generator.manual_seed((int(seed) * 1_000_003 + int(index)) % (2**63 - 1)) + return generator + + +def _rand(generator=None) -> float: + """Draw one uniform in ``[0, 1)`` from ``generator``, or from the global RNG.""" + return float(torch.rand(1, generator=generator).item()) + + +def _uniform(low: float, high: float, generator=None) -> float: + """Draw one uniform in ``[low, high)``.""" + return low + (high - low) * _rand(generator) + + +def _resolve_channel(dd, key, default): + """Read a channel index the sample may carry, falling back to a configured default. + + The stack order depends on which channels were requested, so the dataset is the only + thing that reliably knows where the label map sits; a hardcoded index is silently + wrong for some ``input_data_type`` values. + """ + channel = dd.get(key, default) + if torch.is_tensor(channel): + channel = int(channel.item()) + return int(channel) + + +def _shift_volume(volume, delta): + """Translate ``volume`` by ``delta`` voxels along its spatial axes, padding with zeros. + + Deliberately not ``torch.roll``: rolling is circular, so anatomy cut off one face would + reappear on the opposite one and the landmark coordinate would land back *inside* the + volume on the wrong structure - a corruption no bounds check can catch. + + Returns: + ``(shifted, True)``, or ``(volume, False)`` if a shift would empty the volume. + """ + shifted = torch.zeros_like(volume) + src: list = [slice(None)] + dst: list = [slice(None)] + for axis, size in enumerate(volume.shape[1:]): + step = int(delta[axis]) + if abs(step) >= size: + return volume, False + src.append(slice(max(0, -step), size - max(0, step))) + dst.append(slice(max(0, step), size - max(0, -step))) + shifted[tuple(dst)] = volume[tuple(src)] + return shifted, True + + +def _intersect_loss_mask(dd, keep): + """Fold ``keep`` into ``dd['loss_mask']``, creating the key if the sample has none.""" + carried = dd.get("loss_mask") + dd["loss_mask"] = keep if carried is None else (carried.bool() & keep) + + +class LandmarksRandLimitedFov: + """Truncate the volume the way a limited scanner field of view does, then re-centre. + + A cutout is centred offline on the bounding-box centre of its segmentation, and + inference runs the same code - so on a scan whose field of view stops mid-thorax the + box is centred on whatever anatomy is *visible*, and the content sits quite differently + from any untruncated training sample. This reproduces that: zero one or more + axis-aligned slabs, recompute the bounding-box centre of what survives, and shift the + fixed-size box back onto it. Landmarks in the removed slabs are unsupervised. + + The spatial shape never changes, so the model's input size is unaffected. + + Args: + prob: Per spatial axis, the probability that axis is cut at all. Axis order is the + volume's own, i.e. ``dd["input"]`` dim ``axis + 1``. + cut_range: Per axis, the ``(low, high)`` fraction of the axis removed per cut side. + both_sides_prob: Given an axis is cut, the probability both ends are cut rather + than one. + min_keep_fraction: Floor on the surviving extent of an axis, as a fraction. + min_visible_landmarks: Resample the cut if fewer supervised landmarks than this + would survive; after ``max_attempts`` the sample is left untouched. + max_attempts: Bound on that resampling. + recentre: Whether to shift the surviving content back to the volume centre. + label_channel: Channel of ``dd["input"]`` holding the instance-label map, used for + the bounding box. Overridden per sample by ``dd[label_channel_key]``. + label_channel_key: Sample key carrying that channel index. + seed: When set, the cut is a pure function of ``seed`` and ``dd[index_key]``, so a + sample gets the same field of view in every epoch. Leave ``None`` for training. + index_key: Sample key carrying the dataset index. + + Raises: + ValueError: If the ranges are malformed, which would otherwise degrade silently. + """ + + def __init__( + self, + prob=(0.25, 0.10, 0.60), + cut_range=((0.05, 0.30), (0.05, 0.20), (0.10, 0.45)), + both_sides_prob: float = 0.35, + min_keep_fraction: float = 0.25, + min_visible_landmarks: int = 16, + max_attempts: int = 8, + recentre: bool = True, + label_channel: int = -1, + label_channel_key: str = "label_channel", + seed: int | None = None, + index_key: str = "sample_index", + ): + prob = tuple(float(p) for p in prob) + cut_range = tuple((float(low), float(high)) for low, high in cut_range) + if len(prob) != 3 or len(cut_range) != 3: + raise ValueError(f"prob and cut_range need one entry per spatial axis; got {len(prob)} and {len(cut_range)}.") + if any(not 0.0 <= p <= 1.0 for p in prob): + raise ValueError(f"prob entries must lie in [0, 1]; got {prob}.") + if any(not 0.0 <= low <= high < 0.5 for low, high in cut_range): + raise ValueError(f"cut_range entries must satisfy 0 <= low <= high < 0.5; got {cut_range}.") + if not 0.0 <= both_sides_prob <= 1.0: + raise ValueError(f"both_sides_prob must lie in [0, 1]; got {both_sides_prob}.") + if not 0.0 < min_keep_fraction <= 1.0: + raise ValueError(f"min_keep_fraction must lie in (0, 1]; got {min_keep_fraction}.") + if max_attempts < 1: + raise ValueError(f"max_attempts must be at least 1; got {max_attempts}.") + + self.prob = prob + self.cut_range = cut_range + self.both_sides_prob = both_sides_prob + self.min_keep_fraction = min_keep_fraction + self.min_visible_landmarks = min_visible_landmarks + self.max_attempts = max_attempts + self.recentre = recentre + self.label_channel = label_channel + self.label_channel_key = label_channel_key + self.seed = seed + self.index_key = index_key + + def _clamp(self, frac_low: float, frac_high: float) -> tuple[float, float]: + """Scale both cut depths down proportionally until enough of the axis survives.""" + total = frac_low + frac_high + excess = total - (1.0 - self.min_keep_fraction) + if total <= 0.0 or excess <= 0.0: + return frac_low, frac_high + return frac_low - excess * (frac_low / total), frac_high - excess * (frac_high / total) + + def _sample_box(self, shape, generator): + """Sample the slab to keep, or ``None`` when no axis was cut.""" + low = [0, 0, 0] + high = list(shape) + cut_any = False + for axis in range(3): + if _rand(generator) >= self.prob[axis]: + continue + range_low, range_high = self.cut_range[axis] + both = _rand(generator) < self.both_sides_prob + cut_low = both or _rand(generator) < 0.5 + frac_low = _uniform(range_low, range_high, generator) if cut_low else 0.0 + frac_high = _uniform(range_low, range_high, generator) if (both or not cut_low) else 0.0 + frac_low, frac_high = self._clamp(frac_low, frac_high) + size = shape[axis] + new_low = round(frac_low * size) + new_high = size - round(frac_high * size) + if new_high - new_low < 1 or (new_low == 0 and new_high == size): + continue + low[axis], high[axis] = new_low, new_high + cut_any = True + return (low, high) if cut_any else None + + @staticmethod + def _inside(target, low, high): + """Boolean per landmark: does it lie in the half-open keep box.""" + low_t = torch.as_tensor(low, dtype=target.dtype, device=target.device) + high_t = torch.as_tensor(high, dtype=target.dtype, device=target.device) + return torch.all((target >= low_t) & (target < high_t), dim=1) + + def __call__(self, dd: dict) -> dict: + """Cut, re-centre and re-mask ``dd``; return it untouched when nothing was cut.""" + volume = dd["input"] + target = dd["target"] + shape = tuple(volume.shape[1:]) + generator = _sample_generator(self.seed, dd, self.index_key) + carried = dd.get("loss_mask") + + for _ in range(self.max_attempts): + box = self._sample_box(shape, generator) + if box is None: + # No slab sampled. Returning early matters: re-centring here would undo the + # affine's deliberate off-centre translation on every such sample. + return dd + low, high = box + visible = self._inside(target, low, high) + survivors = visible if carried is None else (visible & carried.bool()) + if int(survivors.sum()) >= self.min_visible_landmarks: + break + else: + return dd + + cropped = torch.zeros_like(volume) + keep = (slice(None), slice(low[0], high[0]), slice(low[1], high[1]), slice(low[2], high[2])) + cropped[keep] = volume[keep] + + delta = torch.zeros(3, dtype=torch.long) + if self.recentre: + channel = _resolve_channel(dd, self.label_channel_key, self.label_channel) + # round() because an affine resamples in float: a label of 43 comes back as + # 43.000004 and a bare `> 0` on the raw values would still be right, but the + # bounding box must not pick up interpolation dust from a zero background. + indices = torch.nonzero(cropped[channel].round() > 0) + if indices.numel() == 0: + return dd # the cut removed every label; a sample with no anatomy teaches nothing + centre = (indices.amin(0) + indices.amax(0)) // 2 + delta = torch.as_tensor(shape) // 2 - centre + + shifted, shifted_ok = _shift_volume(cropped, delta) + if not shifted_ok: + return dd + target = target + delta.to(target.dtype) + + limits = torch.as_tensor(shape, dtype=target.dtype, device=target.device) - 1 + visible = visible & torch.all((target >= 0) & (target <= limits), dim=1) + + dd["input"] = shifted + dd["target"] = target + _intersect_loss_mask(dd, visible) + return dd + + +class LandmarksRandLabelDropout: + """Delete whole instance labels from the input and unsupervise their landmarks. + + A segmentation that misses a structure entirely is a common upstream failure, and it + clusters at the ends of an ordered label range - the first and last ribs, the topmost + and bottommost vertebrae - rather than falling uniformly, which is what ``end_bias`` + encodes. + + Args: + prob: Probability of dropping anything at all. + n_labels: Inclusive ``(min, max)`` number of labels to drop. Never all of them. + contiguous: Drop a run of neighbouring labels rather than a random subset. + end_bias: Probability that a contiguous run starts at one end of the label range. + label_channel: Channel holding the instance-label map; overridden per sample by + ``dd[label_channel_key]``. + label_channel_key: Sample key carrying that channel index. + landmark_labels_key: Sample key mapping each landmark to the label it sits on. When + absent, the volume is still edited but no landmark is unsupervised - so the + dataset must either supply it or re-derive the mask itself. + seed: When set, the draw is a pure function of ``seed`` and ``dd[index_key]``. + index_key: Sample key carrying the dataset index. + + Raises: + ValueError: If ``prob``, ``end_bias`` or ``n_labels`` are malformed. + """ + + def __init__( + self, + prob: float = 0.3, + n_labels=(1, 3), + contiguous: bool = True, + end_bias: float = 0.7, + label_channel: int = -1, + label_channel_key: str = "label_channel", + landmark_labels_key: str = "landmark_labels", + seed: int | None = None, + index_key: str = "sample_index", + ): + n_labels = (int(n_labels[0]), int(n_labels[1])) + if not 0.0 <= prob <= 1.0: + raise ValueError(f"prob must lie in [0, 1]; got {prob}.") + if not 0.0 <= end_bias <= 1.0: + raise ValueError(f"end_bias must lie in [0, 1]; got {end_bias}.") + if not 1 <= n_labels[0] <= n_labels[1]: + raise ValueError(f"n_labels must satisfy 1 <= min <= max; got {n_labels}.") + + self.prob = prob + self.n_labels = n_labels + self.contiguous = contiguous + self.end_bias = end_bias + self.label_channel = label_channel + self.label_channel_key = label_channel_key + self.landmark_labels_key = landmark_labels_key + self.seed = seed + self.index_key = index_key + + def _choose(self, present, generator): + """Pick which of the ``present`` labels to delete, as a list of label values.""" + n_present = int(present.numel()) + # Never delete every label: an input with no anatomy has nothing to learn from. + high = min(self.n_labels[1], n_present - 1) + if high < 1: + return [] + low = min(self.n_labels[0], high) + count = max(low, min(high, int(_uniform(low, high + 1, generator)))) + + if not self.contiguous: + order = torch.randperm(n_present, generator=generator) + return present[order[:count]].tolist() + + if _rand(generator) < self.end_bias: + start = 0 if _rand(generator) < 0.5 else n_present - count + else: + start = int(_uniform(0, n_present - count + 1, generator)) + start = max(0, min(start, n_present - count)) + return present[start : start + count].tolist() + + def __call__(self, dd: dict) -> dict: + """Delete a run of labels from every channel of ``dd['input']``, with probability ``prob``.""" + generator = _sample_generator(self.seed, dd, self.index_key) + if _rand(generator) >= self.prob: + return dd + + volume = dd["input"] + channel = _resolve_channel(dd, self.label_channel_key, self.label_channel) + labels = volume[channel].round() + present = torch.unique(labels) + present = present[present > 0].sort().values + if present.numel() == 0: + return dd + + chosen = self._choose(present, generator) + if not chosen: + return dd + + chosen_t = torch.as_tensor(chosen, dtype=labels.dtype, device=labels.device) + drop = torch.isin(labels, chosen_t) + if not bool(drop.any()): + return dd + # Every channel, not just the label map: a field of view that lost a rib lost its + # surface and its intensities too. + dd["input"] = volume.masked_fill(drop.unsqueeze(0), 0.0) + + landmark_labels = dd.get(self.landmark_labels_key) + if landmark_labels is not None: + gone = torch.isin(landmark_labels, chosen_t.to(landmark_labels.dtype)) + _intersect_loss_mask(dd, ~gone) + return dd + + +class RandLabelCollapse: + """Replace instance labels with one value, so identity cannot be read off the input. + + A label map fed to the network as raw values hands it the instance identity directly. + Collapsing it to a binary mask forces the model to localise from geometry instead, + which is what it has to do whenever the upstream segmentation mislabels. + + This edits ``dd["input"]`` only and belongs in an image-only pipeline, run *after* any + bookkeeping copy of the label map has been split off - collapsing the map a loss mask + is derived from would unsupervise almost everything. + + Args: + prob: Probability of collapsing a given sample. + value: The value every non-zero label becomes. + channels: Channels of ``dd["input"]`` to collapse. When ``None`` the sample must + carry ``dd[channels_key]``. + channels_key: Sample key listing those channels. + seed: When set, the draw is a pure function of ``seed`` and ``dd[index_key]``. + index_key: Sample key carrying the dataset index. + + Raises: + ValueError: If ``prob`` is not a probability. + """ + + def __init__( + self, + prob: float = 0.3, + value: float = 1.0, + channels=None, + channels_key: str = "input_label_channels", + seed: int | None = None, + index_key: str = "sample_index", + ): + if not 0.0 <= prob <= 1.0: + raise ValueError(f"prob must lie in [0, 1]; got {prob}.") + self.prob = prob + self.value = float(value) + self.channels = None if channels is None else [int(c) for c in channels] + self.channels_key = channels_key + self.seed = seed + self.index_key = index_key + + def __call__(self, dd: dict) -> dict: + """Binarise the label channels of ``dd['input']`` with probability ``prob``. + + Raises: + KeyError: If no channels were configured and the sample names none, which would + otherwise make the transform a silent no-op for a whole training run. + """ + generator = _sample_generator(self.seed, dd, self.index_key) + if _rand(generator) >= self.prob: + return dd + + channels = self.channels if self.channels is not None else dd.get(self.channels_key) + if channels is None: + raise KeyError( + f"RandLabelCollapse needs the label channels: pass channels=... or have the dataset set dd[{self.channels_key!r}]." + ) + if torch.is_tensor(channels): + channels = channels.tolist() + + volume = dd["input"] + collapsed = volume.clone() + for channel in channels: + index = int(channel) + collapsed[index] = (volume[index] > 0).to(volume.dtype) * self.value + dd["input"] = collapsed + return dd + + class Compose: def __init__(self, transforms): # `None` entries are dropped rather than rejected: callers compose an @@ -217,15 +634,52 @@ def __call__(self, dd): return dd -def create_transform(config): - transform_type = config["type"] +#: Config ``"type"`` string -> transform class. A project that needs its own variants +#: builds a new dict from this one and passes it to the factories below; see +#: ``rib_poi.data.transforms.RIB_TRANSFORMS``. +TRANSFORMS = { + "LandmarksRandAffine": LandmarksRandAffine, + "LandMarksRandHorizontalFlip": LandMarksRandHorizontalFlip, + "LandMarksRandHorizontalFlipNeighbor": LandMarksRandHorizontalFlipNeighbor, + "LandmarksRandLimitedFov": LandmarksRandLimitedFov, + "LandmarksRandLabelDropout": LandmarksRandLabelDropout, + "RandLabelCollapse": RandLabelCollapse, +} + + +def create_transform(config, registry=None): + """Build one transform from a ``{"type", "params"}`` config. - if transform_type == "LandmarksRandAffine": - return LandmarksRandAffine(**config["params"]) - elif transform_type == "LandMarksRandHorizontalFlip": - return LandMarksRandHorizontalFlip(**config["params"]) - else: - raise ValueError(f"Unknown transform type: {transform_type}") + Args: + config: The config mapping. + registry: Name-to-class mapping to resolve against; :data:`TRANSFORMS` by default. + + Returns: + The constructed transform. + + Raises: + UnknownTypeError: If the config names a transform that is not registered. The + previous hand-written factory raised a bare ``ValueError`` and could only + reach two of the classes in this module. + """ + return build(TRANSFORMS if registry is None else registry, "transform", config) + + +def create_transforms(config, registry=None): + """Build a transform pipeline from ``None``, one config, or a list of configs. + + Args: + config: ``None``, a single ``{"type", "params"}`` mapping, or a list of them. + registry: Name-to-class mapping to resolve against; :data:`TRANSFORMS` by default. + + Returns: + A list of transforms, or ``None`` when there is nothing to build - so the result + can be passed straight to a dataset's ``transforms=`` argument either way. + """ + if config is None: + return None + configs = [config] if isinstance(config, dict) else list(config) + return [create_transform(entry, registry) for entry in configs] or None """ diff --git a/verpex/models/densenet.py b/verpex/models/densenet.py index a7f49a4..7fc8f20 100644 --- a/verpex/models/densenet.py +++ b/verpex/models/densenet.py @@ -316,17 +316,19 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Te normalized_heatmaps = torch.softmax(normalized_heatmaps, dim=-1) normalized_heatmaps = normalized_heatmaps.view(B, N, *spatial_shape) - # Apply the heatmaps to the feature maps to get the feature encoding for each landmark - # Expand feature map dimensions from (B, F, H, W, D) to (B, 1, F, H, W, D) for broadcasting - feature_map_expanded = feature_map.unsqueeze(1) - - # Perform weighted sum - # heatmaps_normalized is (B, N, H, W, D) - # feature_map_expanded is (B, 1, F, H, W, D) - # We want the output to be (B, N, F), summing over the spatial dimensions (H, W, D) + # Apply the heatmaps to the feature maps to get the feature encoding for each landmark: + # a (B, N, F) heatmap-weighted sum of the feature map over the spatial dimensions. if self.weight_features: - # weight the feature maps with the heatmaps - landmark_features = (normalized_heatmaps.unsqueeze(2).detach() * feature_map_expanded).sum(dim=(3, 4, 5)) + # Contract straight to (B, N, F). Broadcasting (B, N, 1, H, W, D) against + # (B, 1, F, H, W, D) and summing computes the same thing, but materialises a + # (B, N, F, H, W, D) intermediate first - 20 MB per sample at N=35 and a + # 8x8x9 heatmap, but 5.6 GB at N=104 and 28x36x52, which is what makes larger + # landmark sets and the U-Net backbone run out of memory. + landmark_features = torch.einsum( + "bnv,bfv->bnf", + normalized_heatmaps.detach().flatten(2), + feature_map.flatten(2), + ) else: # global average pooling of feature maps (no heatmap weighting) @@ -449,9 +451,14 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Te B, N, *spatial = heatmaps.shape normalized_heatmaps = torch.softmax(heatmaps.view(B, N, -1), dim=-1).view(B, N, *spatial) - feature_map_expanded = feature_map.unsqueeze(1) if self.weight_features: - landmark_features = (normalized_heatmaps.unsqueeze(2).detach() * feature_map_expanded).sum(dim=(3, 4, 5)) + # See HeatmapDenseNet.forward: einsum avoids a (B, N, F, H, W, D) intermediate, + # which at this decoder's resolution would dominate the memory budget. + landmark_features = torch.einsum( + "bnv,bfv->bnf", + normalized_heatmaps.detach().flatten(2), + feature_map.flatten(2), + ) else: global_features = feature_map.mean(dim=(2, 3, 4)) landmark_features = global_features.unsqueeze(1).expand(-1, N, -1) diff --git a/verpex/modules/data_modules.py b/verpex/modules/data_modules.py index 94c81fa..8673c11 100644 --- a/verpex/modules/data_modules.py +++ b/verpex/modules/data_modules.py @@ -28,7 +28,7 @@ ) from verpex.data.dataloading import get_files as _get_files from verpex.data.dataset import PoiDataset, PoiNeighborDataset, SpineDataset, SpineNeighborDataset -from verpex.data.transforms import create_transform +from verpex.data.transforms import create_transforms from verpex.registry import build @@ -84,7 +84,7 @@ def __init__( input_shape: tuple = (128, 128, 96), zoom: tuple = (1, 1, 1), flip_prob: float = 0.5, - transform_config: dict | None = None, + transform_config: dict | list[dict] | None = None, include_com: bool = False, include_poi_list=None, include_vert_list=None, @@ -178,7 +178,7 @@ def setup(self, stage=None) -> None: # noqa: ARG002 - Lightning passes it; unus # Bound unconditionally: the datasets below always take a `transforms=` # argument, and `transform_config=None` is the declared default. - transform = [create_transform(self.transform_config)] if self.transform_config is not None else None + transform = create_transforms(self.transform_config) if self.dataset in SINGLE_VERTEBRA_DATASETS: self.train_dataset = SpineDataset( @@ -345,7 +345,7 @@ def __init__( input_shape: tuple = (128, 128, 96), zoom: tuple = (1, 1, 1), flip_prob: float = 0.5, - transform_config: dict | None = None, + transform_config: dict | list[dict] | None = None, include_com: bool = False, include_poi_list=None, include_vert_list=None, diff --git a/verpex/modules/poi_module.py b/verpex/modules/poi_module.py index 5e9e3b8..922c2ba 100644 --- a/verpex/modules/poi_module.py +++ b/verpex/modules/poi_module.py @@ -60,6 +60,12 @@ class PoiPredictionModule(pl.LightningModule): freeze_feature_extractor(): Freezes the feature extraction module. """ + #: Registries the two stages are built from. A subclass targeting different anatomy + #: can point these at its own registries and inherit everything else; without them the + #: only way in is to reimplement __init__. + COARSE_MODULES = FEATURE_EXTRACTION_MODULES + REFINER_MODULES = REFINEMENT_MODULES + def __init__( self, coarse_config, @@ -74,8 +80,8 @@ def __init__( super().__init__() if loss_weights is None: loss_weights = [1, 1] - self.feature_extraction_module = create_feature_extraction_module(coarse_config) - self.refinement_module = create_refinement_module(refinement_config) + self.feature_extraction_module = build(self.COARSE_MODULES, "feature extraction module", coarse_config) + self.refinement_module = build(self.REFINER_MODULES, "refinement module", refinement_config) self.lr = lr # Per-submodule LRs (fall back to the module-level lr). Honouring these # lets configs use a smaller LR for the transformer refiner (standard for