Affine augmentations for 3D tensors targeted at medical images

transforms.py

import math
from typing import Union, NamedTuple

import numpy as np
from scipy.ndimage import affine_transform


class AffineTransform(NamedTuple):
    matrix: np.ndarray
    offset: np.ndarray

    def __add__(self, other):
        return type(self)(matrix=self.matrix @ other.matrix,
                          offset=other.matrix @ self.offset + other.offset)

    def __call__(self, image: np.ndarray, order=1, cval=-1000, output_shape=None, **opts):
        t = self
        if image.ndim > t.ndim:
            t = self.extend(image.ndim)
        return affine_transform(image, matrix=t.matrix, offset=t.offset,
                                output_shape=output_shape, order=order, cval=cval, **opts)

    @property
    def ndim(self):
        return len(self.matrix)

    def extend(self, ndim):
        if ndim > self.ndim:
            m = np.eye(ndim, dtype='float64')
            o = np.zeros(ndim, dtype='float64')
            s = ndim - self.ndim
            m[s:, s:] = self.matrix
            o[s:] = self.offset
            return type(self)(matrix=m, offset=o)
        return self

    @classmethod
    def identity(cls, dims=3):
        return cls(matrix=np.eye(dims, dtype='float64'), offset=np.zeros(dims, dtype='float64'))

    def __repr__(self):
        def prepend_lines(text, prefix):
            return ('\n' + prefix).join(text.split('\n'))

        matrix_repr = prepend_lines(str(self.matrix), '\t       ')
        offset_repr = str(self.offset)
        return f"{type(self).__name__}(\n\tmatrix={matrix_repr},\n\toffset={offset_repr}\n)"

    __str__ = __repr__


def rotate(input_shape, angle, axis) -> AffineTransform:
    """ rotate 3d tensor along given axis in 3D """
    angle = np.pi / 180 * angle
    cos = math.cos(angle)
    sin = math.sin(angle)

    iz, iy, ix = input_shape[-3:]
    offset = np.zeros((3,), dtype='float64')

    if axis == 0:
        matrix = np.array([[1, 0, 0],
                           [0, cos, sin],
                           [0, -sin, cos]], dtype='float64')
    elif axis == 1:
        matrix = np.array([[cos, 0, -sin],
                           [0, 1, 0],
                           [sin, 0, cos]], dtype='float64')
    elif axis == 2:
        matrix = np.array([[cos, sin, 0],
                           [-sin, cos, 0],
                           [0, 0, 1]], dtype='float64')
    else:
        raise ValueError()

    offset[0] = iz / 2.0 - 0.5
    offset[1] = iy / 2.0 - 0.5
    offset[2] = ix / 2.0 - 0.5
    offset = np.dot(matrix, offset)
    tmp = np.zeros((3,), dtype='float64')
    tmp[0] = iz / 2.0 - 0.5
    tmp[1] = iy / 2.0 - 0.5
    tmp[2] = ix / 2.0 - 0.5
    offset = tmp - offset
    return AffineTransform(matrix=matrix, offset=offset).extend(len(input_shape))


def scale(input_shape, factor, center=True) -> AffineTransform:
    """ center scale """
    f = np.asarray(factor)
    assert f.shape[0] == 3
    s = np.asarray(input_shape)
    ndim = s.shape[0]
    z = np.ones(ndim)
    z[-3:] = f
    z = 1 / z
    m = np.eye(ndim) * z
    if center:
        o = (s - m @ s) / 2
    else:
        o = np.zeros(len(input_shape))
    return AffineTransform(matrix=m, offset=o)


def resize(image_shape, target_shape):
    return scale(image_shape, np.asarray(target_shape) / image_shape, center=False)


def crop_target_context(*, image_shape, image_scale, store_context, target_context) -> (AffineTransform, np.array):
    crop_pix_each_size = (np.asarray(store_context) - np.asarray(target_context)) / np.asarray(image_scale)
    orig_shape = np.asarray(image_shape)
    small_shape = orig_shape - crop_pix_each_size * 2
    trf = scale(image_shape, orig_shape / small_shape)
    return trf, small_shape


def resize_and_crop(*, image_shape, target_shape, image_scale, store_context, target_context) -> AffineTransform:
    context_crop_trf, context_crop_shape = crop_target_context(image_shape=image_shape, image_scale=image_scale,
                                                               store_context=store_context,
                                                               target_context=target_context)
    final_trf = context_crop_trf + scale(image_shape, np.asarray(target_shape) / context_crop_shape, center=False)
    return final_trf