fix(impact-reg): masked features, MIND statistics and a worker-safe sentinel

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)

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)
 

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.

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 {