File size: 3,324 Bytes
5652fbb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | 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)
|