Add padding of images.

spatial_transformation/pad_image.py

+"""This module provides functionality to pad and center images."""
+
+from typing import Tuple
+
+import numpy as np
+
+
+def extend_image_with_affine(
+        image: np.ndarray, affine: np.ndarray,
+        target_shape: Tuple[int]) -> Tuple[np.ndarray, np.ndarray]:
+    """Extend image boundaries with affine to a target shape.
+
+    Args:
+        image (np.ndarray): [description]
+        affine (np.ndarray): [description]
+        target_shape (Tuple[int]): [description]
+
+    Returns:
+        Tuple[np.ndarray, np.ndarray]: [description]
+    """
+    image_shape = image.shape
+    size_diff = np.subtract(target_shape, image_shape)
+
+    if np.any(size_diff < 0):
+        raise ValueError('Target shape is smaller than image shape.')
+
+    size_diff_even = np.mod(size_diff, 2)
+
+    image_extended = np.zeros(target_shape, dtype=np.float32)
+
+    start = ((size_diff - size_diff_even) / 2).astype(np.int32)
+    end = start + image_shape
+
+    image_extended[start[0]:end[0], start[1]:end[1], start[2]:end[2]] = image
+    affine_translation = np.matmul(affine[:3, :3], start)
+
+    affine_extended = affine
+    affine_extended[:3, 3] -= affine_translation
+
+    return image_extended, affine_extended
+
+
+def pad_to_target_shape(image: np.ndarray,
+                        target_shape: Tuple[int]) -> np.ndarray:
+    """Pad an image to a target shape.
+
+    Args:
+        image (np.ndarray): 3D image
+        target_shape (Tuple[int]): target shape
+
+    Returns:
+        np.ndarray: padded image
+    """
+    size_diff = np.subtract(target_shape, image.shape)
+    size_diff = np.where(size_diff < 0, 0, size_diff)
+
+    pad_widths_before = ((size_diff - np.mod(size_diff, 2)) / 2).astype(
+        np.int32)
+    pad_widths_after = size_diff - pad_widths_before
+
+    pad_width = tuple(zip(pad_widths_before, pad_widths_after))
+
+    return np.pad(image, pad_width)