diff --git a/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py b/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py index c14024a9..5017dd5d 100644 --- a/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py +++ b/apps/impact_reg/impact_reg_konfai/models/elastix_engine.py @@ -217,8 +217,17 @@ def _apply_map_overrides( replaced = " ".join(per_token[token_key] for _ in values.split()) line = f"{indent}({key} {replaced})" lines.append(line) - # Overrides never inject keys, so a knob set for a key absent from every map silently does nothing: # - # surface it (e.g. final_grid_spacing on a rigid-only preset). + # A raw ``parameter_overrides`` entry is the escape hatch for ANY elastix parameter, including one + # the preset's map never mentions (an ITK default such as RequiredRatioOfValidSamples). Replacing + # only pre-existing keys made those silently do nothing, so an absent exact override is APPENDED. + # The named knobs (final_grid_spacing, spatial_samples, ...) keep replacing what exists and only + # warn: injecting them into a map that does not use them is meaningless (e.g. a B-spline grid + # spacing in a rigid-only preset). + missing_exact = [(key, value) for key, value in exact if key not in seen] + if missing_exact: + lines.append("// appended by impact_reg_konfai parameter_overrides") + lines += [f"({key} {value})" for key, value in missing_exact] + seen.update(key for key, _ in missing_exact) for key in sorted(requested - seen): print(f"[ImpactReg] note: override '{key}' matched no entry in the preset's parameter maps.") return "\n".join(lines) @@ -422,8 +431,6 @@ def forward( moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) - dvf_np = self._engine.register( - fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img - ) + dvf_np = self._engine.register(fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img) combined.append(torch.from_numpy(dvf_np)) return torch.stack(combined, dim=0).to(fixed.device) diff --git a/apps/impact_reg/impact_reg_konfai/models/fireants.py b/apps/impact_reg/impact_reg_konfai/models/fireants.py index 5e2c439d..05dc737a 100644 --- a/apps/impact_reg/impact_reg_konfai/models/fireants.py +++ b/apps/impact_reg/impact_reg_konfai/models/fireants.py @@ -389,6 +389,20 @@ def __init__(self, ref: str, in_channels: int, weights: list[float], distance: s self.model_path = hf_hub_download(repo, filename, repo_type="model") # nosec B615 self.model = None # lazy-loaded on the first forward, like IMPACTReg + def preprocessing(self, tensor: torch.Tensor, attribute: list) -> list[torch.Tensor]: + """KonfAI's preprocessing, with the intensity statistics FLATTENED for a single image. + + ``IMPACTReg.preprocessing`` emits ``stats`` as ``[B, 4]``. The MIND TorchScript model branches on + ``stats.numel() == 4`` and then indexes ``stats[0]`` expecting a SCALAR minimum, so a ``[1, 4]`` + tensor makes it subtract a 4-vector from the volume ("size of tensor a (32) must match the size + of tensor b (4)"). The C++ itk-impact metric passes the four values flat, which is what the model + was traced against; FireANTs always registers one pair at a time, so flatten that case here. + """ + prepared = super().preprocessing(tensor, attribute) + if prepared[2].shape[0] == 1: + prepared[2] = prepared[2].reshape(-1) + return prepared + @staticmethod def _stats(tensor: torch.Tensor) -> dict: detached = tensor.detach() @@ -403,8 +417,16 @@ def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor: # if self.model is None: self.model = torch.jit.load(self.model_path) # nosec B614 self.model.to(moved.device).eval() + # FireANTs' masked mode carries the mask as one extra trailing channel on BOTH images (see + # ``apply_mask_to_image``). The feature model wants the image alone, so split the mask off and + # hand it to KonfAI's masked feature loss (nearest-resampled onto every feature layer): the + # metric is then evaluated inside the fixed mask, which is what an elastix/ITK mask means too. + mask: torch.Tensor | None = None + if moved.shape[1] == self.in_channels + 1 and fixed.shape[1] == self.in_channels + 1: + mask = (fixed[:, -1:] > 0.5).to(torch.uint8) + moved, fixed = moved[:, :-1], fixed[:, :-1] with _no_texpr_fuser(): - loss, true_nb = self._compute(moved, [self._stats(moved)], fixed, [self._stats(fixed)], None) + loss, true_nb = self._compute(moved, [self._stats(moved)], fixed, [self._stats(fixed)], mask) return loss / max(true_nb, 1) @@ -435,6 +457,25 @@ def forward(self, moved: torch.Tensor, fixed: torch.Tensor) -> torch.Tensor: return total +def _mask_on_grid(mask: "sitk.Image", image: "sitk.Image", device: str): + """A FireANTs ``Image`` of ``mask`` carrying ``image``'s geometry. + + The mask is defined on the image's grid, but its header can differ by float rounding once it has + crossed KonfAI's Attribute round-trip; FireANTs' ``concatenate`` then rejects it with a spurious + "different physical spaces" (surfacing as a ``TypeError`` in its ``check_and_raise_cond``). Reusing + the image's spacing/origin/direction makes the two spaces identical by construction. + """ + from fireants.io import Image + + return Image( + mask, + device=device, + spacing=image.GetSpacing(), + origin=image.GetOrigin(), + direction=image.GetDirection(), + ) + + class FireANTsEngine: """Register a fixed/moving pair with FireANTs (Rigid -> Affine -> [SyN | Greedy | none]); return the displacement field on the fixed grid. @@ -605,8 +646,14 @@ def register( use_moving_mask = self._is_partial_mask(moving_mask) masked = use_fixed_mask or use_moving_mask if masked: - fmask = Image(fixed_mask, device=device) if use_fixed_mask else generate_image_mask_allones(fixed_img) - mmask = Image(moving_mask, device=device) if use_moving_mask else generate_image_mask_allones(moving_img) + fmask = ( + _mask_on_grid(fixed_mask, fixed, device) if use_fixed_mask else generate_image_mask_allones(fixed_img) + ) + mmask = ( + _mask_on_grid(moving_mask, moving, device) + if use_moving_mask + else generate_image_mask_allones(moving_img) + ) fixed_img = apply_mask_to_image(fixed_img, fmask) moving_img = apply_mask_to_image(moving_img, mmask) @@ -766,9 +813,7 @@ def forward( moving_img = data_to_image(moving[b].detach().cpu().numpy(), moving_attrs[b]) fixed_mask_img = data_to_image(fixed_mask[b].detach().cpu().numpy(), fmask_attrs[b]) moving_mask_img = data_to_image(moving_mask[b].detach().cpu().numpy(), mmask_attrs[b]) - dvf_np = self._engine.register( - fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img - ) + dvf_np = self._engine.register(fixed_img, moving_img, device_index, fixed_mask_img, moving_mask_img) combined.append(torch.from_numpy(dvf_np)) return torch.stack(combined, dim=0).to(fixed.device) diff --git a/konfai/data/patching.py b/konfai/data/patching.py index 6adc7084..bbb3bac9 100644 --- a/konfai/data/patching.py +++ b/konfai/data/patching.py @@ -93,8 +93,23 @@ # VOLUME it allows is what DatasetManager._sweep_tile then shapes into the block actually read. SWEEP_SLAB_ROWS = 64 + # "not looked up yet", where None is itself an answer (a store with no read granularity to state). -_UNRESOLVED = object() +# A class with a by-name ``__reduce__`` rather than a bare ``object()``: the manager is pickled into +# every DataLoader worker (spawn), and a plain ``object()`` unpickles as a NEW instance, so the +# ``is _UNRESOLVED`` test failed in the worker and the sentinel itself got indexed +# (``'object' object is not subscriptable`` in ``_sweep_rows``). +class _Unresolved: + __slots__ = () + + def __reduce__(self) -> str: + return "_UNRESOLVED" + + def __repr__(self) -> str: + return "_UNRESOLVED" + + +_UNRESOLVED = _Unresolved() # The bytes each element travels as through a sweep (float32). What a sweep holds in those elements # is _sweep_resident_regions, and DatasetManager.sweep_block_bytes prices it. diff --git a/konfai/utils/ome_zarr.py b/konfai/utils/ome_zarr.py index 6d162bc5..94363cf4 100644 --- a/konfai/utils/ome_zarr.py +++ b/konfai/utils/ome_zarr.py @@ -39,6 +39,7 @@ import itertools import operator import shutil +import tempfile import threading from collections import OrderedDict from collections.abc import Sequence @@ -48,6 +49,7 @@ import numpy as np +from konfai.utils import uri from konfai.utils.errors import DatasetManagerError from konfai.utils.runtime import map_over_rank_pool @@ -177,6 +179,22 @@ def _konfai_attributes(store_path: str) -> dict[str, Any]: return {} +def _from_ngff_zarr(store_path: str | Path) -> Any: + """ngff-zarr's multiscales for ``store_path``. A remote root goes in as a key-to-bytes mapping + over its own filesystem (ngff-zarr >= 0.44 reads a remote string as a local path) and as its URL + when ngff-zarr refuses the mapping, which older releases resolve themselves.""" + if not uri.is_uri(store_path): + return ngff_zarr.from_ngff_zarr(str(store_path)) + # Through uri.filesystem, so a missing fsspec backend or configuration is the structured + # DatasetManagerError, never a raw dependency error. + filesystem = uri.filesystem(store_path) + _, target = uri.split_scheme(str(store_path)) + try: + return ngff_zarr.from_ngff_zarr(filesystem.get_mapper(target)) + except ValueError: + return ngff_zarr.from_ngff_zarr(str(store_path)) + + @lru_cache(maxsize=8) def _load_image(store_path: str, level: int) -> Any: """Return the ``NgffImage`` for ``level`` of an OME-Zarr store, memoised per (store, level). @@ -194,7 +212,7 @@ def _load_image(store_path: str, level: int) -> Any: """ _require_ngff_zarr() try: - multiscales = ngff_zarr.from_ngff_zarr(str(store_path)) + multiscales = _from_ngff_zarr(store_path) except (KeyError, IndexError, OSError, TypeError, ValueError) as exc: raise DatasetManagerError( f"Cannot open OME-Zarr store '{store_path}' (level {level}).", @@ -258,7 +276,7 @@ def is_displacement_field(store_path: str | Path) -> bool: if not _NGFF_ZARR_AVAILABLE: return False try: - metadata = ngff_zarr.from_ngff_zarr(str(store_path)).metadata + metadata = _from_ngff_zarr(store_path).metadata except Exception: # "Not a displacement field" is the only answer this owes: it is asked purely to decide HOW to # read an entry, and an absent or unreadable store is not one either. @@ -608,7 +626,7 @@ def place(coords: tuple) -> None: def _level_path(store_path: str, level: int) -> str | None: """The zarr path of one level, from the store's multiscales metadata, memoised beside the image.""" try: - datasets = ngff_zarr.from_ngff_zarr(store_path).metadata.datasets + datasets = _from_ngff_zarr(store_path).metadata.datasets return str(datasets[level if len(datasets) > 1 else 0].path) except Exception: return None @@ -624,7 +642,32 @@ def _read_level_window(store_path: str, level: int, image: Any, index: tuple) -> array = _level_array(store_path, level_path) if tuple(array.shape) == tuple(image.data.shape): return _read_chunked(store_path, level_path, array, index) - return np.asarray(image.data[index]) + return _lazy_window(image.data, index) + + +def _lazy_window(data: Any, index: tuple) -> np.ndarray: + """The plain lazy read. A lazy array that refuses a stepped slice (ngff-zarr >= 0.44 wraps the + level in an adapter that takes unit steps only) is read over the unit-step span and stepped here.""" + try: + return np.asarray(data[index]) + except NotImplementedError: + # Each slice normalized against the shape, then its ascending unit-step span; a negative + # step reads that span backwards from its own end, which lands on the indices the original + # slice named. + bounds = tuple( + slice(*item.indices(size)) if isinstance(item, slice) else item + for item, size in zip(index, data.shape, strict=True) + ) + span = tuple( + slice(item.stop + 1, item.start + 1) + if isinstance(item, slice) and item.step < 0 + else slice(item.start, item.stop) + if isinstance(item, slice) + else item + for item in bounds + ) + steps = tuple(slice(None, None, item.step) for item in bounds if isinstance(item, slice)) + return np.asarray(data[span])[steps] def _store_index( @@ -813,6 +856,11 @@ def write_ome_zarr( exist only from 0.6, so a caller passing both could only ever pass them consistently: an invariant worth removing rather than documenting. """ + if scale_factors and uri.is_uri(store_path): + raise DatasetManagerError( + f"Cannot append pyramid levels to the remote store '{store_path}'.", + "Levels are derived in place through local paths; write the store locally and upload it.", + ) array_data = np.asarray(data) # The one write path: the store described and created empty (ngff-zarr's metadata, the # caller's chunking), filled by zarr itself, its levels grafted beside level 0. Handing @@ -876,6 +924,30 @@ def _type_component_axis(multiscales: Any, axis_type: str) -> None: axis.type = axis_type +def _write_skeleton(store_path: str | Path, multiscales: Any, version: str) -> None: + """ngff-zarr's metadata for the store, written in place. ngff-zarr (>= 0.44) writes local + directories only, so a remote root gets the skeleton written locally and uploaded through the + root's own filesystem: a few bytes of metadata, before the array is created underneath it.""" + if not uri.is_uri(store_path): + ngff_zarr.to_ngff_zarr(str(store_path), multiscales, overwrite=True, version=version) + return + filesystem = uri.filesystem(store_path) + _, target = uri.split_scheme(str(store_path)) + if "/" not in target.strip("/"): + raise DatasetManagerError( + f"Refusing to create the store at the filesystem root '{store_path}'.", + "Name a key under the root (e.g. '.../dataset.ome.zarr'): creating a store replaces what its path holds.", + ) + if filesystem.exists(target): + filesystem.rm(target, recursive=True) + filesystem.makedirs(target, exist_ok=True) + with tempfile.TemporaryDirectory() as scratch: + local = Path(scratch) / "skeleton" + ngff_zarr.to_ngff_zarr(str(local), multiscales, overwrite=True, version=version) + for file in sorted(path for path in local.rglob("*") if path.is_file()): + filesystem.put_file(str(file), uri.join(target, file.relative_to(local).as_posix())) + + def create_ome_zarr_store( store_path: str | Path, shape: Sequence[int], @@ -934,7 +1006,7 @@ def create_ome_zarr_store( _type_component_axis(multiscales, _DISPLACEMENT_AXIS_TYPE) version = _RFC5_VERSION # version is explicit because to_ngff_zarr defaults to 0.5, which zarr-python 2 cannot write. - ngff_zarr.to_ngff_zarr(str(store_path), multiscales, overwrite=True, version=version) + _write_skeleton(store_path, multiscales, version) # The level-0 key comes from the metadata rather than a literal: ngff-zarr builds it from the # image name, so "scale0/image" is its convention to change, not ours to hardcode. @@ -1077,13 +1149,18 @@ def append_ome_zarr_levels( KonfAI attribute sidecar is untouched, being a key beside theirs; a displacement field keeps its typed component axis through the same call that types it at creation. """ + if uri.is_uri(store_path): + raise DatasetManagerError( + f"Cannot append pyramid levels to the remote store '{store_path}'.", + "Levels are derived in place through local paths; write the store locally and upload it.", + ) _require_ngff_zarr() if not scale_factors: return store = Path(store_path) clear_ome_zarr_cache(store) field = is_displacement_field(store) - base = ngff_zarr.from_ngff_zarr(str(store)).images[0] + base = _from_ngff_zarr(store).images[0] stored_chunks = tuple(int(size) for size in base.data.chunksize) if downsample_method in (None, "ITKWASM_BIN_SHRINK"): multiscales = _bin_shrink_multiscales(base, _level_zero_scale_factors(scale_factors), stored_chunks) @@ -1152,7 +1229,7 @@ def get_ome_zarr_info(store_path: str | Path, level: int = 0) -> dict[str, Any]: image = _load_image(str(store_path), level) dims = [str(axis).lower() for axis in image.dims] try: - n_levels = len(ngff_zarr.from_ngff_zarr(str(store_path)).images) + n_levels = len(_from_ngff_zarr(store_path).images) except (OSError, TypeError, ValueError): n_levels = 1 return {