Skip to content

Commit

Permalink
Add random bias field transform
Browse files Browse the repository at this point in the history
  • Loading branch information
fepegar committed Jan 1, 2020
1 parent e5f5971 commit 26ebc05
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 0 deletions.
1 change: 1 addition & 0 deletions torchio/transforms/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from .random_noise import RandomNoise
from .random_affine import RandomAffine
from .random_motion import RandomMotion
from .random_bias_field import RandomBiasField
from .random_elastic_deformation import RandomElasticDeformation
from .interpolation import Interpolation

Expand Down
83 changes: 83 additions & 0 deletions torchio/transforms/random_bias_field.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
"""
Adapted from NiftyNet
"""

import numpy as np
import torch
from ..torchio import INTENSITY
from ..utils import is_image_dict
from .random_transform import RandomTransform


class RandomBiasField(RandomTransform):
def __init__(
self,
coefficients_range=(-0.5, 0.5),
order=3,
seed=None,
verbose=False,
):
super().__init__(seed=seed, verbose=verbose)
self.coefficients_range = coefficients_range
self.order = order

def apply_transform(self, sample):
coefficients = self.get_params(self.order, self.coefficients_range)
sample['random_bias_field'] = coefficients
for image_name, image_dict in sample.items():
if not is_image_dict(image_dict):
continue
if image_dict['type'] != INTENSITY:
continue
coefficients = self.get_params(self.order, self.coefficients_range)
sample[image_name]['random_bias_field'] = coefficients
image_dict['data'] *= self.generate_bias_field_map(
image_dict['data'], self.order, coefficients)
return sample

@staticmethod
def get_params(order, coefficients_range):
"""
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 z_order in range(0, order + 1 - (x_order + y_order)):
n = torch.FloatTensor(1).uniform_(*coefficients_range)
random_coefficients.append(n.item())
return np.array(random_coefficients)

@staticmethod
def generate_bias_field_map(data, order, coefficients):
"""
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) for n in half_shape]

bf_map = np.zeros(shape)
x_mesh, y_mesh, z_mesh = np.asarray(np.meshgrid(*ranges))

x_mesh /= x_mesh.max()
y_mesh /= y_mesh.max()
z_mesh /= z_mesh.max()

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)):
random_coefficient = coefficients[i]
new_map = (
random_coefficient
* x_mesh ** x_order
* y_mesh ** y_order
* z_mesh ** z_order
)
bf_map += np.transpose(new_map, (1, 0, 2)) # why?
i += 1
return np.exp(bf_map)

0 comments on commit 26ebc05

Please sign in to comment.