wip

Author: dkuegler

FastSurferCNN/data_loader/conform.py

@@ -28,7 +28,7 @@
 from numpy import typing as npt
 
 from FastSurferCNN.utils import AffineMatrix4x4, ScalarType, Shape1d, logging, nibabelHeader, nibabelImage
-from FastSurferCNN.utils.affines import aff2axcodes, OrntArrayType
+from FastSurferCNN.utils.affines import OrntArrayType, aff2axcodes, io_orientation
 from FastSurferCNN.utils.arg_types import ImageSizeOption, OrientationType, StrictOrientationType, VoxSizeOption
 from FastSurferCNN.utils.arg_types import float_gt_zero_and_le_one as __conform_to_one_mm
 from FastSurferCNN.utils.arg_types import img_size as __img_size
@@ -55,7 +55,7 @@
 Date: May-12-2025
 """
 FIX_MGH_AFFINE_CALCULATION = False
-FIX_CENTER_NOT_CENTER = False
+FIX_CENTER_NOT_CENTER = True
 
 LOGGER = logging.getLogger(__name__)
 
@@ -290,7 +290,7 @@ def do_nothing(data: _TB) -> _TB:
         return image_data, do_nothing
 
     # vox2vox should always be a "soft" transform
-    vox2vox = affine_for_target_orientation(source_affine, "soft " + _target_orientation, image_data.shape).round()
+    vox2vox = vox2vox_for_target_orientation(source_affine, "soft " + _target_orientation, image_data.shape).round()
     if np.allclose(vox2vox, np.eye(4)):  # is already target_orientation
         return image_data, do_nothing
     else:  # is not target_affine yet
@@ -300,7 +300,7 @@ def do_nothing(data: _TB) -> _TB:
         return apply_vox2vox(image_data, vox2vox, out_shape), inverse_apply_vox2vox
 
 
-def affine_for_target_orientation(
+def vox2vox_for_target_orientation(
         source_affine: AffineMatrix4x4,
         target_orientation: OrientationType,
         shape: npt.ArrayLike,
@@ -956,11 +956,12 @@ def _validate(param_name, param_type, param_value):
     source_vox_size = img.header.get_zooms()
 
     source_affine = get_affine_from_any(img)
-    _vox2vox = affine_for_target_orientation(source_affine, orientation, (0,) * 3)[:3, :3]
+    _vox2vox = vox2vox_for_target_orientation(source_affine, orientation, (0,) * 3)[:3, :3]
     _target_affine = source_affine[:3, :3] @ _vox2vox
     h1["Mdc"] = np.linalg.inv(_target_affine / np.linalg.norm(_target_affine, axis=0, keepdims=True))
-    re_order_axes = np.argmax(np.abs(_vox2vox), axis=1).tolist()
-    # re_order_axes = nib.io_orientation(_vox2vox)[:, 0].tolist()
+    _vox2vox_hom = np.eye(4, dtype=_vox2vox.dtype)
+    _vox2vox_hom[:3, :3] = _vox2vox
+    re_order_axes = io_orientation(_vox2vox_hom)[:, 0].astype(int).tolist()
 
     shape: list[int] = [(source_img_shape if target_img_size is None else target_img_size)[i] for i in re_order_axes]
     h1.set_data_shape(shape + [1])

FastSurferCNN/run_prediction.py

@@ -39,7 +39,7 @@
 
 import FastSurferCNN.reduce_to_aseg as rta
 from FastSurferCNN.data_loader import data_utils as du
-from FastSurferCNN.data_loader.conform import affine_for_target_orientation, conform, is_conform, to_target_orientation
+from FastSurferCNN.data_loader.conform import vox2vox_for_target_orientation, conform, is_conform, to_target_orientation
 from FastSurferCNN.inference import Inference
 from FastSurferCNN.quick_qc import check_volume
 from FastSurferCNN.utils import PLANES, Plane, logging, nibabelImage, parser_defaults
@@ -388,7 +388,7 @@ def get_prediction(
 
         orig_in_lia, back_to_native = to_target_orientation(orig_data, affine, target_orientation="LIA")
         shape = orig_in_lia.shape + (self.get_num_classes(),)
-        _ornt_transform, _ = affine_for_target_orientation(affine, target_orientation="LIA")
+        _ornt_transform, _ = vox2vox_for_target_orientation(affine, target_orientation="LIA")
         _zoom = _zoom[_ornt_transform[:, 0]]
 
         pred_prob = torch.zeros(shape, **kwargs)

test/image/conftest.py

@@ -27,7 +27,7 @@ def vox_size(request) -> float:
     return request.param
 
 
-@pytest.fixture(scope="session", params=[4, 6]) # [128, 256])
+@pytest.fixture(scope="session", params=[8, 15]) # [128, 256])
 def img_size(request) -> float:
     return request.param
 

test/image/test_conform_reorient.py

@@ -7,55 +7,57 @@
 from pytest import approx
 
 from FastSurferCNN.data_loader.conform import OrientationType, conform, prepare_mgh_header
-from FastSurferCNN.utils import AffineMatrix4x4, Image3d, nibabelHeader
+from FastSurferCNN.utils import AffineMatrix4x4, Image3d, nibabelHeader, nibabelImage
+from FastSurferCNN.utils.affines import aff2axcodes
 from FastSurferCNN.utils.arg_types import StrictOrientationType
 
 logger = getLogger(__name__)
 
 class MultiCoordImages(TypedDict):
-    X: nib.Nifti1Image
-    Y: nib.Nifti1Image
-    Z: nib.Nifti1Image
+    X: nibabelImage
+    Y: nibabelImage
+    Z: nibabelImage
 
 conform_reorient = {"rescale": None, "dtype": np.float32}
 
 
-def circle_data(img_size: int) -> Image3d[np.float32]:
+def circle_data(img_size: int, radius: float, center: float) -> Image3d[np.float32]:
     """Generates a 3D image with a centered sphere of radius img_size/2."""
-    data = np.mgrid[0:img_size, 0:img_size, 0:img_size].astype(np.float32) - (img_size - 1) / 2.0
-    return (np.sum(data * data, axis=0) < img_size * img_size / 4.0).astype(np.float32)
+    data = np.mgrid[0:img_size, 0:img_size, 0:img_size].astype(np.float32) - center
+    return (np.sum(data * data, axis=0) < radius * radius).astype(np.float32)
 
 
 @pytest.fixture(scope="session")
-def circle_image(random_affine: AffineMatrix4x4, img_size: int) -> nib.Nifti1Image:
-    return nib.Nifti1Image(circle_data(img_size), random_affine)
-
+def radius(img_size: int) -> float:
+    return img_size / 2.0 - 2.0
 
 @pytest.fixture(scope="session")
-def random_image(random_affine: AffineMatrix4x4, img_size: int) -> nib.Nifti1Image:
-    return nib.Nifti1Image(np.random.randn(img_size, img_size, img_size), random_affine)
+def center(img_size: int) -> float:
+    return (img_size - 1) / 2.0
 
 
 @pytest.fixture(scope="session")
-def empty_image(random_affine: AffineMatrix4x4, img_size: int) -> nib.Nifti1Image:
-    return nib.Nifti1Image(np.empty((img_size,) * 3), random_affine)
+def circle_image(random_affine: AffineMatrix4x4, img_size: int, radius: float, center: float) -> nib.Nifti1Image:
+    return nib.Nifti1Image(circle_data(img_size, radius, center), random_affine)
+
+
+@pytest.fixture(scope="session", params=[1.0, 0.23])
+def resample_factor(request) -> float:
+    return request.param
 
 
 @pytest.fixture(scope="session")
-def worldcoord_images(random_affine: AffineMatrix4x4, img_size: int) -> MultiCoordImages:
-    data = worldcoords_data(random_affine, img_size)
-    return MultiCoordImages(**{c: nib.Nifti1Image(data[..., i], random_affine) for i, c in enumerate("XYZ")})
+def random_image(random_affine: AffineMatrix4x4, img_size: int) -> nib.Nifti1Image:
+    return nib.Nifti1Image(np.random.randn(img_size, img_size, img_size), random_affine)
 
 
-def worldcoords_data(affine: AffineMatrix4x4, img_size: int) -> np.ndarray:
-    xi = np.moveaxis(np.mgrid[0:img_size, 0:img_size, 0:img_size], 0, -1)
-    return nib.affines.apply_affine(affine, xi.reshape((-1, 3)).astype(float)).reshape(xi.shape).astype(np.float32)
+@pytest.fixture(scope="session")
+def empty_image(random_affine: AffineMatrix4x4, img_size: int) -> nib.Nifti1Image:
+    return nib.Nifti1Image(np.empty((img_size,) * 3), random_affine)
 
 
 def affine2orientation(affine: AffineMatrix4x4) -> OrientationType:
     """Generates the orientation type string from an affine matrix."""
-    from FastSurferCNN.utils.affines import aff2axcodes
-
     orientation: StrictOrientationType = "".join(aff2axcodes(affine, ("LR", "PA", "IS")))
     # make sure the affine is normalized for vox_sizes not 1
     norm_affine = affine[:3, :3] / np.linalg.norm(affine[:3, :3], keepdims=True, axis=0)
@@ -121,17 +123,20 @@ def image(
             soft_orientation: OrientationType,
             img_size: int,
             vox_size: float,
+            resample_factor: float,
     ) -> nib.MGHImage:
         """
         Conform `circle_image` to `soft_orientation` and back to the original orientation.
         """
-        from FastSurferCNN.utils.affines import aff2axcodes
-
-        there = conform(circle_image, orientation=soft_orientation, **conform_reorient)
+        there = conform(
+            circle_image,
+            orientation=soft_orientation, vox_size=vox_size * resample_factor, img_size=64, order=1,
+            **conform_reorient,
+        )
         in_orientation: OrientationType = "soft " + "".join(aff2axcodes(circle_image.affine, ("LR", "PA", "IS")))
-        return conform(there, orientation=in_orientation, **conform_reorient, img_size=img_size, vox_size=vox_size)
+        return conform(there, orientation=in_orientation, img_size=img_size, vox_size=vox_size, **conform_reorient)
 
-    def test_affine(self, image: nib.MGHImage, circle_image: nib.Nifti1Image) -> None:
+    def test_affine(self, image: nib.MGHImage, circle_image: nibabelImage) -> None:
         """
         Tests whether the affines of the original and the re-oriented images are the same.
         """
@@ -140,57 +145,88 @@ def test_affine(self, image: nib.MGHImage, circle_image: nib.Nifti1Image) -> Non
         # currently, the translation parts of the affines differ
         assert actual == approx(expected, abs=1e-5), "The affines of original and re-reoriented images differ!"
 
-    def test_image(self, image: nib.MGHImage, circle_image: nib.Nifti1Image):
+    def test_image(
+            self,
+            image: nib.MGHImage,
+            circle_image: nibabelImage,
+            radius: float,
+            center: float,
+            soft_orientation: AffineMatrix4x4,
+    ) -> None:
         """
         Tests whether the content of the original and the re-oriented images are the same in the "center circle".
         """
 
         # this has to be filtered by the region of the image that is the same
-        inner_circle = circle_data(circle_image.shape[0] - 2)
-        outer_circle = circle_data(circle_image.shape[0] + 2)[1:-1, 1:-1, 1:-1]
-        mask = np.pad(inner_circle, 1, constant_values=0) + (1 - outer_circle)
-        expected = np.where(mask, circle_image.dataobj, np.nan)
+        inner_circle = circle_data(circle_image.shape[0], radius - 1, center)
+        outer_circle = circle_data(circle_image.shape[0], radius + 1, center)
+        # the expected image is 0 outside the outer circle, 1 inside the inner circle and ignored (NaN) in between,
+        # where interpolation can cause differences
+        expected = np.where(outer_circle - inner_circle, np.nan, circle_image.dataobj)
         actual = np.asarray(image.dataobj)
-        assert actual == approx(expected, abs=0.3, nan_ok=True), "The data differs from the re-oriented image!"
-
+        logger.info(f"Difference is {np.max(np.abs(actual - circle_image.get_fdata()))} - {np.nanmax(np.abs(actual - expected))}")
 
-def test_reorient_worldcoords(
-        worldcoord_images: MultiCoordImages,
-        soft_orientation: OrientationType,
-        vox_size: float,
-        img_size: int,
-):
-    """
-    This test checks, whether the world coordinates are consistent.
-
-    In effect this means generating world coordinates for an original image (xyz) and
+        assert actual == approx(expected, abs=0.3, nan_ok=True), "The data differs from the re-oriented image!"
 
-    Parameters
-    ----------
-    worldcoord_images : MultiCoordImages
-        A dictionary of three world coordinate images (X, Y, Z) generated based on random_affine.
-    soft_orientation : OrientationType
-        The soft orientation to which the world coordinate images will be conformed.
+class TestReorientWorldCoords:
 
+    @staticmethod
+    def worldcoords_data(affine: AffineMatrix4x4, img_size: int) -> np.ndarray:
+        xi = np.moveaxis(np.mgrid[0:img_size, 0:img_size, 0:img_size], 0, -1)
+        return nib.affines.apply_affine(affine, xi.reshape((-1, 3)).astype(float)).reshape(xi.shape).astype(np.float32)
 
-    Returns
-    -------
+    @pytest.fixture(scope="class")
+    def worldcoord_images(self, random_affine: AffineMatrix4x4, img_size: int) -> MultiCoordImages:
+        data = self.worldcoords_data(random_affine, img_size)
+        return MultiCoordImages(**{c: nib.Nifti1Image(data[..., i], random_affine) for i, c in enumerate("XYZ")})
 
-    """
-    from pytest import approx, xfail
+    @pytest.fixture(scope="class")
+    def conf_img_size(self) -> int:
+        return 16
 
-    from FastSurferCNN.data_loader.conform import conform
+    @pytest.fixture(scope="class")
+    def conf_images(
+            self,
+            conf_img_size: int,
+            worldcoord_images: MultiCoordImages,
+            orientation: OrientationType,
+    ) -> MultiCoordImages:
+        return MultiCoordImages(
+            **{k: conform(img, **conform_reorient, img_size=conf_img_size) for k, img in worldcoord_images.items()},
+        )
+
+    @pytest.mark.parametrize(argnames=["dim_name"], argvalues=[["X"], ["Y"], ["Z"]])
+    def test_reorient_worldcoords_affine(
+            self,
+            conf_images: MultiCoordImages,
+            orientation: OrientationType,
+            dim_name: str,
+    ):
+        """
+        This test checks, whether the affines of the world coordinate images are consistent with the original affine.
+        """
+        actual = "".join(aff2axcodes(conf_images[dim_name], ("LR", "PA", "IS")))
+        expected = orientation
+        assert actual == expected
 
-    xfail("This test is currently broken!")
 
-    conf_img_size = 8
+    def test_reorient_worldcoords_image(
+            self,
+            conf_images: MultiCoordImages,
+            soft_orientation: OrientationType,
+            vox_size: float,
+            img_size: int,
+    ):
+        """
+        This test checks, whether the world coordinates are consistent.
+        """
+        from pytest import approx, xfail
 
-    ximg, yimg, zimg = [conform(img, **conform_reorient, img_size=conf_img_size) for img in worldcoord_images.values()]
-    xyz = np.stack([ximg.get_fdata(), yimg.get_fdata(), zimg.get_fdata()], axis=-1)
-    logger.info("Checking affines of world images:")
+        xyz = np.stack([img.get_fdata() for img in conf_images], axis=-1)
+        logger.info("Checking affines of world images:")
 
-    worldcoords = worldcoords_data(ximg.affine, ximg.shape[0])
-    center_mask = np.pad(circle_data(ximg.shape[0] - 2), 1, constant_values=0)
-    expected = np.where(center_mask[..., None].astype(np.bool_), worldcoords, np.NaN)
+        worldcoords = self.worldcoords_data(conf_images[0].affine, conf_images[0].shape[0])
+        # center_mask = np.pad(circle_data(conf_images[0] - 2), 1, constant_values=0)
+        expected = np.where(center_mask[..., None].astype(np.bool_), worldcoords, np.NaN)
 
-    assert xyz == approx(expected, abs=0.1), "The world images "
+        assert xyz == approx(expected, abs=0.1), "The world images "

test/image/test_orientation_transform.py

@@ -7,7 +7,7 @@
 from pytest import approx
 
 from FastSurferCNN.data_loader.conform import (
-    affine_for_target_orientation,
+    vox2vox_for_target_orientation,
     does_vox2vox_rot_require_interpolation,
     ornt2affine,
 )
@@ -87,7 +87,7 @@ def test_ornt2affine_data(strict_orientation: OrientationType, img_size: int):
 
 def test_affine_for_target_orientation(random_affine: AffineMatrix4x4, img_size: int, orientation: OrientationType):
     """Test whether affine_for_target_orientation works as expected."""
-    vox2vox = affine_for_target_orientation(random_affine, orientation, (img_size,) * 3)
+    vox2vox = vox2vox_for_target_orientation(random_affine, orientation, (img_size,) * 3)
     combined = random_affine @ vox2vox
     actual = "".join(aff2axcodes(combined, ("lr", "pa", "is")))
     expected = orientation.lower().removeprefix("soft").lstrip("-_ ")
@@ -116,7 +116,7 @@ def test_affine_for_target_orientation2(
         strict_orientation: StrictOrientationType,
 ):
     """Test whether affine_for_target_orientation works as expected."""
-    vox2vox = affine_for_target_orientation(random_affine, "soft " + strict_orientation, (img_size,) * 3)
+    vox2vox = vox2vox_for_target_orientation(random_affine, "soft " + strict_orientation, (img_size,) * 3)
     is_reorder_flip = not does_vox2vox_rot_require_interpolation(vox2vox)
     assert is_reorder_flip, "affine_for_target_orientation did not yield a \"soft\" vox2vox transformation!"