from collections import defaultdict
from typing import Dict
from typing import List
from typing import Tuple
from typing import Union
import numpy as np
import torch
from .. import RandomTransform
from ... import IntensityTransform
from ....data.subject import Subject
from ....typing import TypeData
[docs]
class RandomBiasField(RandomTransform, IntensityTransform):
r"""Add random MRI bias field artifact.
MRI magnetic field inhomogeneity creates intensity
variations of very low frequency across the whole image.
The bias field is modeled as a linear combination of
polynomial basis functions, as in K. Van Leemput et al., 1999,
*Automated model-based tissue classification of MR images of the brain*.
It was implemented in NiftyNet by Carole Sudre and used in
`Sudre et al., 2017, Longitudinal segmentation of age-related
white matter hyperintensities
<https://www.sciencedirect.com/science/article/pii/S1361841517300257?via%3Dihub>`_.
Args:
coefficients: Maximum magnitude :math:`n` of polynomial coefficients.
If a tuple :math:`(a, b)` is specified, then
:math:`n \sim \mathcal{U}(a, b)`.
order: Order of the basis polynomial functions.
**kwargs: See :class:`~torchio.transforms.Transform` for additional
keyword arguments.
"""
def __init__(
self,
coefficients: Union[float, Tuple[float, float]] = 0.5,
order: int = 3,
**kwargs,
):
super().__init__(**kwargs)
self.coefficients_range = self._parse_range(
coefficients,
'coefficients_range',
)
self.order = _parse_order(order)
def apply_transform(self, subject: Subject) -> Subject:
arguments: Dict[str, dict] = defaultdict(dict)
for image_name in self.get_images_dict(subject):
coefficients = self.get_params(
self.order,
self.coefficients_range,
)
arguments['coefficients'][image_name] = coefficients
arguments['order'][image_name] = self.order
transform = BiasField(**self.add_include_exclude(arguments))
transformed = transform(subject)
assert isinstance(transformed, Subject)
return transformed
def get_params(
self,
order: int,
coefficients_range: Tuple[float, float],
) -> List[float]:
# Sampling of the appropriate number of coefficients for the creation
# of the bias field map
random_coefficients = []
for x_order in range(0, order + 1):
for y_order in range(0, order + 1 - x_order):
for _ in range(0, order + 1 - (x_order + y_order)):
sample = self.sample_uniform(*coefficients_range)
random_coefficients.append(sample)
return random_coefficients
class BiasField(IntensityTransform):
r"""Add MRI bias field artifact.
Args:
coefficients: Magnitudes of the polinomial coefficients.
order: Order of the basis polynomial functions.
**kwargs: See :class:`~torchio.transforms.Transform` for additional
keyword arguments.
"""
def __init__(
self,
coefficients: Union[List[float], Dict[str, List[float]]],
order: Union[int, Dict[str, int]],
**kwargs,
):
super().__init__(**kwargs)
self.coefficients = coefficients
self.order = order
self.invert_transform = False
self.args_names = ['coefficients', 'order']
def arguments_are_dict(self):
coefficients_dict = isinstance(self.coefficients, dict)
order_dict = isinstance(self.order, dict)
if coefficients_dict != order_dict:
message = 'If one of the arguments is a dict, all must be'
raise ValueError(message)
return coefficients_dict and order_dict
def apply_transform(self, subject: Subject) -> Subject:
coefficients, order = self.coefficients, self.order
for name, image in self.get_images_dict(subject).items():
if self.arguments_are_dict():
assert isinstance(self.coefficients, dict)
assert isinstance(self.order, dict)
coefficients, order = self.coefficients[name], self.order[name]
assert isinstance(order, int)
bias_field = self.generate_bias_field(
image.data,
order,
coefficients, # type: ignore[arg-type]
)
if self.invert_transform:
np.divide(1, bias_field, out=bias_field)
image.set_data(image.data * torch.as_tensor(bias_field))
return subject
@staticmethod
def generate_bias_field(
data: TypeData,
order: int,
coefficients: TypeData,
) -> np.ndarray:
# Create the bias field map using a linear combination of polynomial
# functions and the coefficients previously sampled
shape = np.array(data.shape[1:]) # first axis is channels
half_shape = shape / 2
ranges = [np.arange(-n, n) + 0.5 for n in half_shape]
bias_field = np.zeros(shape)
meshes = np.asarray(np.meshgrid(*ranges))
for mesh in meshes:
mesh_max = mesh.max()
if mesh_max > 0:
mesh /= mesh_max
x_mesh, y_mesh, z_mesh = meshes
i = 0
for x_order in range(order + 1):
for y_order in range(order + 1 - x_order):
for z_order in range(order + 1 - (x_order + y_order)):
coefficient = coefficients[i]
new_map = (
coefficient
* x_mesh**x_order
* y_mesh**y_order
* z_mesh**z_order
)
bias_field += np.transpose(new_map, (1, 0, 2)) # why?
i += 1
bias_field = np.exp(bias_field).astype(np.float32)
return bias_field
def _parse_order(order):
if not isinstance(order, int):
raise TypeError(f'Order must be an int, not {type(order)}')
if order < 0:
raise ValueError(f'Order must be a positive int, not {order}')
return order