Finalize raw prediction artifact stage

Author: Donglai Wei

connectomics/data/processing/__init__.py

@@ -1,16 +1,16 @@
 # Core processing functions
 from .bbox import bbox_ND, crop_ND, replace_ND
 from .bbox_processor import BBoxInstanceProcessor, BBoxProcessorConfig
-from .iou import seg_to_iou, segs_to_iou
 
 # Pipeline builder (primary entry point for label transforms)
 from .build import create_label_transform_pipeline
+from .iou import seg_to_iou, segs_to_iou
 
 # Utility functions used by decoding
 from .misc import get_seg_type
 
 # MONAI-native transforms and composition
-from .nnunet_preprocess import NNUNetPreprocessd
+from .nnunet_preprocess import NNUNetPreprocessd, restore_prediction_to_input_space
 from .transforms import (
     ComputeBinaryRatioWeightd,
     ComputeUNet3DWeightd,
@@ -61,6 +61,7 @@
     "SegSelectiond",
     "MultiTaskLabelTransformd",
     "NNUNetPreprocessd",
+    "restore_prediction_to_input_space",
     # Pipelines
     "create_label_transform_pipeline",
     # IoU

connectomics/data/processing/nnunet_preprocess.py

@@ -22,6 +22,135 @@
 logger = logging.getLogger(__name__)
 
 
+def _infer_spatial_dims_from_array(array: np.ndarray) -> int:
+    if array.ndim <= 2:
+        return array.ndim
+    if array.ndim == 3:
+        return 3
+    return array.ndim - 1
+
+
+def _spatial_shape(array: np.ndarray, spatial_dims: int) -> tuple:
+    if array.ndim == spatial_dims:
+        return tuple(int(v) for v in array.shape)
+    return tuple(int(v) for v in array.shape[-spatial_dims:])
+
+
+def _resample_array_to_shape(
+    array: np.ndarray,
+    target_shape: Sequence[int],
+    spatial_dims: int,
+    order: int,
+) -> np.ndarray:
+    from scipy.ndimage import zoom
+
+    target = tuple(int(v) for v in target_shape)
+    if _spatial_shape(array, spatial_dims) == target:
+        return array
+
+    def _zoom_single(vol: np.ndarray) -> np.ndarray:
+        factors = np.array(target, dtype=np.float32) / np.maximum(
+            np.array(vol.shape, dtype=np.float32), 1.0
+        )
+        return zoom(
+            vol.astype(np.float32, copy=False),
+            zoom=factors,
+            order=order,
+            mode="nearest",
+            prefilter=order > 1,
+        )
+
+    if array.ndim == spatial_dims + 1:
+        channels = [_zoom_single(array[c])[None] for c in range(array.shape[0])]
+        return np.vstack(channels).astype(array.dtype, copy=False)
+    if array.ndim == spatial_dims:
+        return _zoom_single(array).astype(array.dtype, copy=False)
+    return array
+
+
+def _fit_array_to_shape(
+    array: np.ndarray, target_shape: Sequence[int], spatial_dims: int
+) -> np.ndarray:
+    target = tuple(int(v) for v in target_shape)
+    if _spatial_shape(array, spatial_dims) == target:
+        return array
+
+    if array.ndim == spatial_dims + 1:
+        out = np.zeros((array.shape[0], *target), dtype=array.dtype)
+        in_shape = array.shape[1:]
+        write_shape = tuple(min(int(in_shape[d]), target[d]) for d in range(spatial_dims))
+        out_slices = (slice(None),) + tuple(slice(0, w) for w in write_shape)
+        in_slices = (slice(None),) + tuple(slice(0, w) for w in write_shape)
+        out[out_slices] = array[in_slices]
+        return out
+
+    if array.ndim == spatial_dims:
+        out = np.zeros(target, dtype=array.dtype)
+        in_shape = array.shape
+        write_shape = tuple(min(int(in_shape[d]), target[d]) for d in range(spatial_dims))
+        out_slices = tuple(slice(0, w) for w in write_shape)
+        in_slices = tuple(slice(0, w) for w in write_shape)
+        out[out_slices] = array[in_slices]
+        return out
+
+    return array
+
+
+def restore_prediction_to_input_space(sample: np.ndarray, meta: Dict[str, Any]) -> np.ndarray:
+    """Invert nnU-Net preprocessing metadata for one prediction sample."""
+    preprocess_meta = meta.get("nnunet_preprocess")
+    if not isinstance(preprocess_meta, dict) or not preprocess_meta.get("enabled", False):
+        return sample
+
+    array = sample
+    spatial_dims = int(preprocess_meta.get("spatial_dims", _infer_spatial_dims_from_array(array)))
+    is_integer = np.issubdtype(array.dtype, np.integer)
+    interp_order = 0 if is_integer else 1
+
+    if preprocess_meta.get("applied_resample", False):
+        cropped_shape = preprocess_meta.get("cropped_spatial_shape")
+        if isinstance(cropped_shape, (list, tuple)) and len(cropped_shape) == spatial_dims:
+            array = _resample_array_to_shape(
+                array,
+                target_shape=cropped_shape,
+                spatial_dims=spatial_dims,
+                order=interp_order,
+            )
+
+    if preprocess_meta.get("applied_crop", False):
+        bbox = preprocess_meta.get("crop_bbox")
+        original_shape = preprocess_meta.get("original_spatial_shape")
+        if (
+            isinstance(bbox, (list, tuple))
+            and isinstance(original_shape, (list, tuple))
+            and len(bbox) == spatial_dims
+            and len(original_shape) == spatial_dims
+        ):
+            crop_target_shape = tuple(int(b[1]) - int(b[0]) for b in bbox)
+            array = _fit_array_to_shape(array, crop_target_shape, spatial_dims=spatial_dims)
+
+            if array.ndim == spatial_dims + 1:
+                restored = np.zeros((array.shape[0], *original_shape), dtype=array.dtype)
+                slices = tuple(slice(int(b[0]), int(b[1])) for b in bbox)
+                restored[(slice(None), *slices)] = array
+            else:
+                restored = np.zeros(tuple(int(v) for v in original_shape), dtype=array.dtype)
+                slices = tuple(slice(int(b[0]), int(b[1])) for b in bbox)
+                restored[slices] = array
+            array = restored
+
+    transpose_axes = preprocess_meta.get("transpose_axes")
+    if isinstance(transpose_axes, (list, tuple)) and len(transpose_axes) == spatial_dims:
+        inverse_axes = np.argsort(np.asarray(transpose_axes))
+        if array.ndim == spatial_dims + 1:
+            perm = [0] + [int(i) + 1 for i in inverse_axes]
+            array = np.transpose(array, perm)
+        elif array.ndim == spatial_dims:
+            array = np.transpose(array, tuple(int(i) for i in inverse_axes))
+
+    return array
+
+
 class NNUNetPreprocessd(MapTransform):
     """nnU-Net style preprocessing transform."""
 

connectomics/inference/__init__.py

@@ -2,6 +2,7 @@
 
 from .artifact import (
     PredictionArtifactMetadata,
+    build_prediction_artifact_metadata,
     read_prediction_artifact,
     write_prediction_artifact,
     write_prediction_artifact_attrs,
@@ -29,6 +30,7 @@
 __all__ = [
     "InferenceManager",
     "PredictionArtifactMetadata",
+    "build_prediction_artifact_metadata",
     "read_prediction_artifact",
     "write_prediction_artifact",
     "write_prediction_artifact_attrs",

connectomics/inference/artifact.py

@@ -5,7 +5,7 @@
 import json
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
-from typing import Any, Callable, Mapping
+from typing import Any, Callable, Mapping, Sequence
 
 import numpy as np
 
@@ -22,6 +22,10 @@ class PredictionArtifactMetadata:
     input_shape: tuple[int, ...] | None = None
     final_shape: tuple[int, ...] | None = None
     crop_pad: tuple[tuple[int, int], ...] | None = None
+    transpose: tuple[int, ...] | None = None
+    model_architecture: str | None = None
+    model_output_identity: str | None = None
+    decode_after_inference: bool | None = None
     chunk_shape: tuple[int, ...] | None = None
     halo: tuple[int, ...] | None = None
     channel_order: tuple[str, ...] | None = None
@@ -31,6 +35,86 @@ class PredictionArtifactMetadata:
     extra: Mapping[str, Any] = field(default_factory=dict)
 
 
+def _cfg_get(obj: Any, path: str, default: Any = None) -> Any:
+    node = obj
+    for part in path.split("."):
+        if node is None:
+            return default
+        if isinstance(node, Mapping):
+            node = node.get(part, default)
+        else:
+            node = getattr(node, part, default)
+    return node
+
+
+def _tuple_or_none(value: Sequence[Any] | None) -> tuple[int, ...] | None:
+    if value in (None, [], ()):
+        return None
+    return tuple(int(v) for v in value)
+
+
+def _model_output_identity(cfg: Any, output_head: str | None) -> str | None:
+    parts: list[str] = []
+    if output_head:
+        parts.append(f"head={output_head}")
+    else:
+        primary_head = _cfg_get(cfg, "model.primary_head")
+        if primary_head:
+            parts.append(f"primary_head={primary_head}")
+
+    select_channel = _cfg_get(cfg, "inference.select_channel")
+    if select_channel is not None:
+        parts.append(f"select_channel={select_channel}")
+
+    return ";".join(parts) if parts else None
+
+
+def build_prediction_artifact_metadata(
+    cfg: Any,
+    *,
+    image_path: str | None = None,
+    checkpoint_path: str | None = None,
+    output_head: str | None = None,
+    input_shape: Sequence[int] | None = None,
+    final_shape: Sequence[int] | None = None,
+    crop_pad: Sequence[Sequence[int]] | None = None,
+    chunk_shape: Sequence[int] | None = None,
+    halo: Sequence[int] | None = None,
+    intensity_scale: float | None = None,
+    intensity_dtype: str | None = None,
+    extra: Mapping[str, Any] | None = None,
+) -> PredictionArtifactMetadata:
+    """Build standard metadata for raw prediction artifacts."""
+    transform_cfg = _cfg_get(cfg, "inference.prediction_transform")
+    transform_enabled = bool(getattr(transform_cfg, "enabled", False))
+    if intensity_scale is None and transform_enabled:
+        intensity_scale = float(getattr(transform_cfg, "intensity_scale", -1.0))
+    if intensity_dtype is None and transform_enabled:
+        intensity_dtype = getattr(transform_cfg, "intensity_dtype", None)
+
+    return PredictionArtifactMetadata(
+        image_path=image_path,
+        checkpoint_path=str(checkpoint_path) if checkpoint_path is not None else None,
+        output_head=output_head,
+        input_shape=_tuple_or_none(input_shape),
+        final_shape=_tuple_or_none(final_shape),
+        crop_pad=(
+            tuple((int(pair[0]), int(pair[1])) for pair in crop_pad)
+            if crop_pad is not None
+            else None
+        ),
+        transpose=_tuple_or_none(_cfg_get(cfg, "data.data_transform.val_transpose")),
+        model_architecture=_cfg_get(cfg, "model.arch.type"),
+        model_output_identity=_model_output_identity(cfg, output_head),
+        decode_after_inference=bool(_cfg_get(cfg, "inference.decode_after_inference", True)),
+        chunk_shape=_tuple_or_none(chunk_shape),
+        halo=_tuple_or_none(halo),
+        intensity_scale=intensity_scale,
+        intensity_dtype=intensity_dtype,
+        extra=extra or {},
+    )
+
+
 def _json_attr(value: Any) -> Any:
     if value is None or isinstance(value, (str, int, float, bool)):
         return value
@@ -140,6 +224,7 @@ def read_prediction_artifact(
 
 __all__ = [
     "PredictionArtifactMetadata",
+    "build_prediction_artifact_metadata",
     "read_prediction_artifact",
     "write_prediction_artifact",
     "write_prediction_artifact_attrs",

connectomics/inference/chunked.py

@@ -19,7 +19,7 @@
 )
 from ..utils.channel_slices import resolve_channel_indices
 from .artifact import (
-    PredictionArtifactMetadata,
+    build_prediction_artifact_metadata,
     write_prediction_artifact,
 )
 from .lazy import get_lazy_image_reference_shape, lazy_predict_region
@@ -278,7 +278,8 @@ def _run_chunked_prediction_per_rank(
         write_prediction_artifact(
             chunk_path,
             core_pred,
-            metadata=PredictionArtifactMetadata(
+            metadata=build_prediction_artifact_metadata(
+                cfg,
                 image_path=str(image_path),
                 checkpoint_path=str(checkpoint_path) if checkpoint_path is not None else None,
                 output_head=requested_head,
@@ -479,7 +480,8 @@ def write_chunks(dataset) -> None:
 
     write_prediction_artifact(
         output_path,
-        metadata=PredictionArtifactMetadata(
+        metadata=build_prediction_artifact_metadata(
+            cfg,
             image_path=str(image_path),
             checkpoint_path=str(checkpoint_path) if checkpoint_path is not None else None,
             output_head=requested_head,