src/inference/volume_preprocess.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import nibabel as nib
import nibabel.orientations as nio
import numpy as np


@dataclass
class PreprocessContext:
    source_path: Path
    original_img: nib.Nifti1Image
    canonical_img: nib.Nifti1Image
    target_shape: tuple[int, int, int]
    source_slices: tuple[slice, slice, slice]
    dest_slices: tuple[slice, slice, slice]


def _center_crop_or_pad_with_mapping(
    volume: np.ndarray,
    target_shape: tuple[int, int, int],
) -> tuple[np.ndarray, tuple[slice, slice, slice], tuple[slice, slice, slice]]:
    if volume.ndim != 3:
        raise ValueError(f"Expected 3D volume, got shape={volume.shape}")

    in_shape = volume.shape
    out = np.zeros(target_shape, dtype=np.float32)

    in_slices: list[slice] = []
    out_slices: list[slice] = []
    src_slices: list[slice] = []
    dst_slices: list[slice] = []

    for in_len, out_len in zip(in_shape, target_shape):
        if in_len >= out_len:
            start = (in_len - out_len) // 2
            in_slices.append(slice(start, start + out_len))
            out_slices.append(slice(0, out_len))

            # Restore map: target -> canonical.
            src_slices.append(slice(0, out_len))
            dst_slices.append(slice(start, start + out_len))
        else:
            start = (out_len - in_len) // 2
            in_slices.append(slice(0, in_len))
            out_slices.append(slice(start, start + in_len))

            # Restore map: target -> canonical.
            src_slices.append(slice(start, start + in_len))
            dst_slices.append(slice(0, in_len))

    out[tuple(out_slices)] = volume[tuple(in_slices)]
    return out, tuple(src_slices), tuple(dst_slices)


def prepare_image_for_model(
    image_path: Path,
    target_shape: tuple[int, int, int],
) -> tuple[np.ndarray, PreprocessContext]:
    image_path = Path(image_path).expanduser().resolve()
    original_img = nib.load(str(image_path))
    canonical_img = nib.as_closest_canonical(original_img)
    canonical = canonical_img.get_fdata().astype(np.float32)

    prepped, source_slices, dest_slices = _center_crop_or_pad_with_mapping(canonical, target_shape)

    ctx = PreprocessContext(
        source_path=image_path,
        original_img=original_img,
        canonical_img=canonical_img,
        target_shape=target_shape,
        source_slices=source_slices,
        dest_slices=dest_slices,
    )
    return prepped, ctx


def restore_prediction_to_original(
    prediction_in_target_space: np.ndarray,
    ctx: PreprocessContext,
) -> np.ndarray:
    pred = np.asarray(prediction_in_target_space, dtype=np.float32)
    if pred.shape != ctx.target_shape:
        raise ValueError(f"Prediction shape {pred.shape} does not match target shape {ctx.target_shape}")

    canonical_shape = tuple(int(x) for x in ctx.canonical_img.shape[:3])
    pred_canonical = np.zeros(canonical_shape, dtype=np.float32)
    pred_canonical[ctx.dest_slices] = pred[ctx.source_slices]

    can_ornt = nio.io_orientation(ctx.canonical_img.affine)
    orig_ornt = nio.io_orientation(ctx.original_img.affine)
    transform = nio.ornt_transform(can_ornt, orig_ornt)
    pred_original = nio.apply_orientation(pred_canonical, transform)
    return np.asarray(pred_original, dtype=np.float32)