From 26ebc053a9ae3cc60dda70c30f82a072ac0216a7 Mon Sep 17 00:00:00 2001
From: Fernando Perez-Garcia <fepegar@gmail.com>
Date: Wed, 1 Jan 2020 20:13:36 +0000
Subject: [PATCH] Add random bias field transform

---
 torchio/transforms/__init__.py          |  1 +
 torchio/transforms/random_bias_field.py | 83 +++++++++++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 torchio/transforms/random_bias_field.py

diff --git a/torchio/transforms/__init__.py b/torchio/transforms/__init__.py
index c06f55973..7ed9b9577 100644
--- a/torchio/transforms/__init__.py
+++ b/torchio/transforms/__init__.py
@@ -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
 
diff --git a/torchio/transforms/random_bias_field.py b/torchio/transforms/random_bias_field.py
new file mode 100644
index 000000000..4583fc017
--- /dev/null
+++ b/torchio/transforms/random_bias_field.py
@@ -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)