Include cropping functionality.

src/spacial_transformation/crop_image.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+This module provides functionality for image cropping.
+
+Created on Aug 30 18:54:28 2023
+
+@author: jsteiglechner
+"""
+
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+
+
+def crop_image(
+    image: Union[np.ndarray, torch.tensor],
+    start: Tuple[int, int, int],
+    size: Tuple[int, int, int],
+) -> np.ndarray:
+    """
+    Remove unwanted outer areas from an 3D image.
+
+    Parameters
+    ----------
+    image : Union[np.ndarray, torch.tensor]
+        Image to crop.
+    start : Tuple[int, int, int]
+        Start position of crop.
+    size : Tuple[int, int, int]
+        Size of crop.
+
+    Raises
+    ------
+    ValueError
+        If target size is larger than image size.
+
+    Returns
+    -------
+    np.ndarray
+        Cropped image.
+
+    """
+    if size > image.shape:
+        raise ValueError(
+            f'Target size {size} is not allowed to be larger than image {image.shape}.'
+        )
+    end = np.add(start, size)
+
+    return image[start[0]:end[0], start[1]:end[1], start[2]:end[2]]
+
+
+def randomly_crop_two_images(
+        image1: np.ndarray,
+        image2: np.ndarray,
+        size: Tuple[int, int, int],
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Remove same random outer ares from two images.
+
+    Parameters
+    ----------
+    image1 : np.ndarray
+        Image to crop.
+    image2 : np.ndarray
+        Image to crop.
+    size : Tuple[int, int, int]
+        Size of crop.
+
+    Raises
+    ------
+    ValueError
+        If images have different size.
+
+    Returns
+    -------
+    crop1 : np.ndarray
+        Crop of image1.
+    crop2 : np.ndarray
+        Crop of image2.
+
+    """
+    if np.any(image1.shape != image2.shape):
+        raise ValueError('Images do not have the same shape.')
+
+    valid_start_range = np.subtract(image1.shape, size)
+    random_start = tuple(np.random.randint(valid_start_range))
+
+    crop1 = crop_image(image1, random_start, size)
+    crop2 = crop_image(image2, random_start, size)
+
+    return crop1, crop2
+
+
+def gaussian_crop_two_images(
+        image1: np.ndarray,
+        image2: np.ndarray,
+        size: Tuple[int, int, int],
+) -> np.ndarray:
+    """
+    Remove same outer ares with gaussion distribution from two images.
+
+    Parameters
+    ----------
+    image1 : np.ndarray
+        Image to crop.
+    image2 : np.ndarray
+        Image to crop.
+    size : Tuple[int, int, int]
+        Size of crop.
+
+    Raises
+    ------
+    ValueError
+        If images have different size.
+
+    Returns
+    -------
+    crop1 : np.ndarray
+        Crop of image1.
+    crop2 : np.ndarray
+        Crop of image2.
+
+    """
+    if np.any(image1.shape != image2.shape):
+        raise ValueError('Images do not have the same shape.')
+
+    valid_start_range = np.subtract(image1.shape, size)
+    gaussian_start = np.around(
+        np.random.normal(loc=valid_start_range / 2.,
+                         scale=valid_start_range / 4.))
+    gaussian_start = np.where(gaussian_start < 0., 0., gaussian_start)
+    gaussian_start = np.where(gaussian_start > valid_start_range,
+                              valid_start_range, gaussian_start)
+    gaussian_start = tuple(gaussian_start.astype(np.int))
+
+    crop1 = crop_image(image1, gaussian_start, size)
+    crop2 = crop_image(image2, gaussian_start, size)
+
+    return crop1, crop2
+
+
+def crop_background(image: np.ndarray) -> np.ndarray:
+    """
+    Remove all black intersecting planes across entire image.
+
+    Parameters
+    ----------
+    image : np.ndarray
+        Image to crop.
+
+    Returns
+    -------
+    np.ndarray
+        Crop with blank planes removed.
+
+    """
+    mask = ~np.isclose(image, 0.)
+    return image[np.ix_(mask.any(2), mask.any(1), mask.any(0))]
+
+
+def crop_background_only_outside(
+        image: np.ndarray,
+        affine: np.ndarray = None,
+) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+    """
+    Remove all black outer planes while keeping all inner ones.
+
+    Parameters
+    ----------
+    image : np.ndarray
+        Image to crop.
+    affine : np.ndarray, optional
+        Affine matrix of image. The default is None.
+
+    Returns
+    -------
+    Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]
+        Cropped image and optionally its affine.
+
+    """
+    """Remove all black outer planes while keeping all inner ones."""
+    mask = ~np.isclose(image, 0., atol=.1)
+    i, j, k = image.shape
+    mask_i = mask.any(0)
+    mask_j = mask.any(1)
+    mask_k = mask.any(2)
+    i_start, i_end = mask_i.argmax(), i - mask_i[::-1].argmax()
+    j_start, j_end = mask_j.argmax(), j - mask_j[::-1].argmax()
+    k_start, k_end = mask_k.argmax(), k - mask_k[::-1].argmax()
+
+    if affine is None:
+        return image[i_start:i_end, j_start:j_end, k_start:k_end]
+    else:
+        new_affine = affine.copy()
+        new_affine[:, 3] = np.matmul(new_affine,
+                                     [i_start, j_start, k_start, 1.])
+        return image[i_start:i_end, j_start:j_end, k_start:k_end], new_affine
+
+
+def crop_image_with_affine(
+        image: np.ndarray,
+        affine: np.ndarray,
+        target_shape: Tuple[int],
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Crop image centre with affine to a target shape.
+
+    Parameters
+    ----------
+    image : np.ndarray
+        Image to crop.
+    affine : np.ndarray
+        Images corresponding affine.
+    target_shape : Tuple[int]
+        Target size of crop.
+
+    Raises
+    ------
+    ValueError
+        If target size is larger than image size.
+
+    Returns
+    -------
+    image_cropped : np.ndarray
+        Cropped image.
+    affine_cropped : np.ndarray
+        Affine corresponding to cropped image.
+
+    """
+    image_shape = image.shape
+    size_diff = np.subtract(image_shape, target_shape)
+
+    if np.any(size_diff < 0):
+        raise ValueError('Target shape is larger than image shape.')
+
+    size_diff_even = np.mod(size_diff, 2)
+    start = ((size_diff - size_diff_even) / 2).astype(np.int32)
+
+    image_cropped = crop_image(image=image, start=start, size=target_shape)
+    affine_translation = np.matmul(affine[:3, :3], start)
+
+    affine_cropped = affine
+    affine_cropped[:3, 3] += affine_translation
+
+    return image_cropped, affine_cropped