Add 8-fold Test-Time Augmentation to predict, by default off.

Author: tibuch

n2v/models/n2v_standard.py

@@ -22,7 +22,8 @@
 from ..internals.N2V_DataWrapper import N2V_DataWrapper
 from ..internals.n2v_losses import loss_mse, loss_mae
 from ..utils import n2v_utils
-from ..utils.n2v_utils import pm_identity, pm_normal_additive, pm_normal_fitted, pm_normal_withoutCP, pm_uniform_withCP
+from ..utils.n2v_utils import pm_identity, pm_normal_additive, pm_normal_fitted, pm_normal_withoutCP, pm_uniform_withCP, \
+    tta_forward, tta_backward
 from ..nets.unet import build_single_unet_per_channel
 
 from tifffile import imsave
@@ -77,7 +78,7 @@ def __init__(self, config, name=None, basedir='.'):
 
         config is None or isinstance(config, self._config_class) or _raise(
             ValueError("Invalid configuration of type '%s', was expecting type '%s'." % (
-            type(config).__name__, self._config_class.__name__))
+                type(config).__name__, self._config_class.__name__))
         )
         if config is not None and not config.is_valid():
             invalid_attr = config.is_valid(True)[1]
@@ -104,6 +105,7 @@ def __init__(self, config, name=None, basedir='.'):
         if config is None:
             self._find_and_load_weights()
 
+
     def _build(self):
         return self._build_unet(
             n_dim=self.config.n_dim,
@@ -301,7 +303,6 @@ def prepare_for_training(self, optimizer=None, **kwargs):
                 self.callbacks.append(
                     TensorBoard(log_dir=str(self.logdir / 'logs'), write_graph=False, profile_batch=0))
 
-
         if self.config.train_reduce_lr is not None:
             from keras.callbacks import ReduceLROnPlateau
             rlrop_params = self.config.train_reduce_lr
@@ -336,7 +337,7 @@ def __normalize__(self, data, means, stds):
     def __denormalize__(self, data, means, stds):
         return (data * stds) + means
 
-    def predict(self, img, axes, resizer=PadAndCropResizer(), n_tiles=None):
+    def predict(self, img, axes, resizer=PadAndCropResizer(), n_tiles=None, tta=False):
         """
         Apply the network to sofar unseen data. This method expects the raw data, i.e. not normalized.
         During prediction the mean and standard deviation, stored with the model (during data generation), are used
@@ -351,6 +352,8 @@ def predict(self, img, axes, resizer=PadAndCropResizer(), n_tiles=None):
         resizer : class(Resizer), optional(default=PadAndCropResizer())
         n_tiles : tuple(int)
                   Number of tiles to tile the image into, if it is too large for memory.
+        tta     : bool
+                  Use test-time augmentation during prediction.
 
         Returns
         -------
@@ -375,9 +378,17 @@ def predict(self, img, axes, resizer=PadAndCropResizer(), n_tiles=None):
             normalized = self.__normalize__(img[..., np.newaxis], means, stds)
             normalized = normalized[..., 0]
 
-        pred = \
-        self._predict_mean_and_scale(normalized, axes=new_axes, normalizer=None, resizer=resizer, n_tiles=new_n_tiles)[
-            0]
+        if tta:
+            aug = tta_forward(normalized)
+            preds = []
+            for img in aug:
+                preds.append(self._predict_mean_and_scale(img, axes=new_axes, normalizer=None, resizer=resizer,
+                                             n_tiles=new_n_tiles)[0])
+            pred = tta_backward(preds)
+        else:
+            pred = \
+                self._predict_mean_and_scale(normalized, axes=new_axes, normalizer=None, resizer=resizer,
+                                             n_tiles=new_n_tiles)[0]
 
         pred = self.__denormalize__(pred, means, stds)
 
@@ -573,5 +584,3 @@ def get_yml_dict(self, name, description, authors, test_img, axes, patch_shape=N
     @property
     def _config_class(self):
         return N2VConfig
-
-

n2v/utils/n2v_utils.py

@@ -5,14 +5,14 @@
 
 def get_subpatch(patch, coord, local_sub_patch_radius):
     start = np.maximum(0, np.array(coord) - local_sub_patch_radius)
-    end = start + local_sub_patch_radius*2 + 1
+    end = start + local_sub_patch_radius * 2 + 1
 
     shift = np.minimum(0, patch.shape - end)
 
     start += shift
     end += shift
 
-    slices = [ slice(s, e) for s, e in zip(start, end)]
+    slices = [slice(s, e) for s, e in zip(start, end)]
 
     return patch[tuple(slices)]
 
@@ -40,17 +40,19 @@ def normal_withoutCP(patch, coords, dims):
             rand_coords = random_neighbor(patch.shape, coord)
             vals.append(patch[tuple(rand_coords)])
         return vals
+
     return normal_withoutCP
 
 
 def pm_uniform_withCP(local_sub_patch_radius):
     def random_neighbor_withCP_uniform(patch, coords, dims):
         vals = []
         for coord in zip(*coords):
-            sub_patch = get_subpatch(patch, coord,local_sub_patch_radius)
+            sub_patch = get_subpatch(patch, coord, local_sub_patch_radius)
             rand_coords = [np.random.randint(0, s) for s in sub_patch.shape[0:dims]]
             vals.append(sub_patch[tuple(rand_coords)])
         return vals
+
     return random_neighbor_withCP_uniform
 
 
@@ -60,6 +62,7 @@ def pixel_gauss(patch, coords, dims):
         for coord in zip(*coords):
             vals.append(np.random.normal(patch[tuple(coord)], pixel_gauss_sigma))
         return vals
+
     return pixel_gauss
 
 
@@ -71,6 +74,7 @@ def local_gaussian(patch, coords, dims):
             axis = tuple(range(dims))
             vals.append(np.random.normal(np.mean(sub_patch, axis=axis), np.std(sub_patch, axis=axis)))
         return vals
+
     return local_gaussian
 
 
@@ -80,17 +84,18 @@ def identity(patch, coords, dims):
         for coord in zip(*coords):
             vals.append(patch[coord])
         return vals
+
     return identity
 
 
 def manipulate_val_data(X_val, Y_val, perc_pix=0.198, shape=(64, 64), value_manipulation=pm_uniform_withCP(5)):
     dims = len(shape)
     if dims == 2:
-        box_size = np.round(np.sqrt(100/perc_pix)).astype(np.int)
+        box_size = np.round(np.sqrt(100 / perc_pix)).astype(np.int)
         get_stratified_coords = dw.__get_stratified_coords2D__
         rand_float = dw.__rand_float_coords2D__(box_size)
     elif dims == 3:
-        box_size = np.round(np.sqrt(100/perc_pix)).astype(np.int)
+        box_size = np.round(np.sqrt(100 / perc_pix)).astype(np.int)
         get_stratified_coords = dw.__get_stratified_coords3D__
         rand_float = dw.__rand_float_coords3D__(box_size)
 
@@ -99,7 +104,7 @@ def manipulate_val_data(X_val, Y_val, perc_pix=0.198, shape=(64, 64), value_mani
     Y_val *= 0
     for j in tqdm(range(X_val.shape[0]), desc='Preparing validation data: '):
         coords = get_stratified_coords(rand_float, box_size=box_size,
-                                            shape=np.array(X_val.shape)[1:-1])
+                                       shape=np.array(X_val.shape)[1:-1])
         for c in range(n_chan):
             indexing = (j,) + coords + (c,)
             indexing_mask = (j,) + coords + (c + n_chan,)
@@ -112,17 +117,52 @@ def manipulate_val_data(X_val, Y_val, perc_pix=0.198, shape=(64, 64), value_mani
 
 
 def autocorrelation(x):
-  """
-  nD autocorrelation
-  remove mean per-patch (not global GT)
-  normalize stddev to 1
-  value at zero shift normalized to 1...
-  """
-  x = (x - np.mean(x))/np.std(x)
-  x  = np.fft.fftn(x)
-  x  = np.abs(x)**2
-  x = np.fft.ifftn(x).real
-  x = x / x.flat[0]
-  x = np.fft.fftshift(x)
-  return x
-
+    """
+    nD autocorrelation
+    remove mean per-patch (not global GT)
+    normalize stddev to 1
+    value at zero shift normalized to 1...
+    """
+    x = (x - np.mean(x)) / np.std(x)
+    x = np.fft.fftn(x)
+    x = np.abs(x) ** 2
+    x = np.fft.ifftn(x).real
+    x = x / x.flat[0]
+    x = np.fft.fftshift(x)
+    return x
+
+
+def tta_forward(x):
+    """
+    Augments x 8-fold: all 90 deg rotations plus lr flip of the four rotated versions.
+
+    Parameters
+    ----------
+    x: data to augment
+
+    Returns
+    -------
+    Stack of augmented x.
+    """
+    x_aug = [x, np.rot90(x, 1), np.rot90(x, 2), np.rot90(x, 3)]
+    x_aug_flip = x_aug.copy()
+    for x_ in x_aug:
+        x_aug_flip.append(np.fliplr(x_))
+    return x_aug_flip
+
+
+def tta_backward(x_aug):
+    """
+    Inverts `tta_forward` and averages the 8 images.
+
+    Parameters
+    ----------
+    x_aug: stack of 8-fold augmented images.
+
+    Returns
+    -------
+    average of de-augmented x_aug.
+    """
+    x_deaug = [x_aug[0], np.rot90(x_aug[1], -1), np.rot90(x_aug[2], -2), np.rot90(x_aug[3], -3),
+               np.fliplr(x_aug[4]), np.rot90(np.fliplr(x_aug[5]), -1), np.rot90(np.fliplr(x_aug[6]), -2), np.rot90(np.fliplr(x_aug[7]), -3)]
+    return np.mean(x_deaug, 0)

tests/test_n2v_utils.py

@@ -1,44 +1,46 @@
 import numpy as np
 from n2v.utils import n2v_utils
+from n2v.utils.n2v_utils import tta_forward, tta_backward
+
 
 def test_get_subpatch():
     patch = np.arange(100)
-    patch.shape = (10,10)
+    patch.shape = (10, 10)
 
     subpatch_target = np.array([[11, 12, 13, 14, 15],
                                 [21, 22, 23, 24, 25],
                                 [31, 32, 33, 34, 35],
                                 [41, 42, 43, 44, 45],
                                 [51, 52, 53, 54, 55]])
 
-    subpatch_test = n2v_utils.get_subpatch(patch, (3,3), 2)
+    subpatch_test = n2v_utils.get_subpatch(patch, (3, 3), 2)
 
     assert np.sum(subpatch_target - subpatch_test) == 0
 
-    subpatch_test = n2v_utils.get_subpatch(patch, (3,3), 1)
+    subpatch_test = n2v_utils.get_subpatch(patch, (3, 3), 1)
 
     assert np.sum(subpatch_target[1:-1, 1:-1] - subpatch_test) == 0
 
     patch = np.arange(1000)
-    patch.shape = (10,10,10)
+    patch.shape = (10, 10, 10)
 
-    subpatch_target = np.array([[[31,32,33],
-                                 [41,42,43],
-                                 [51,52,53]],
-                                [[131,132,133],
-                                 [141,142,143],
-                                 [151,152,153]],
-                                [[231,232,233],
-                                 [241,242,243],
-                                 [251,252,253]]])
+    subpatch_target = np.array([[[31, 32, 33],
+                                 [41, 42, 43],
+                                 [51, 52, 53]],
+                                [[131, 132, 133],
+                                 [141, 142, 143],
+                                 [151, 152, 153]],
+                                [[231, 232, 233],
+                                 [241, 242, 243],
+                                 [251, 252, 253]]])
 
-    subpatch_test = n2v_utils.get_subpatch(patch, (1,4,2), 1)
+    subpatch_test = n2v_utils.get_subpatch(patch, (1, 4, 2), 1)
 
     assert np.sum(subpatch_target - subpatch_test) == 0
 
 
 def test_random_neighbor():
-    coord = np.array([51,52,32])
+    coord = np.array([51, 52, 32])
     shape = [128, 128, 128]
 
     for i in range(1000):
@@ -54,9 +56,9 @@ def test_random_neighbor():
 
 def test_pm_normal_neighbor_withoutCP():
     patch = np.arange(100)
-    patch.shape = (10,10)
+    patch.shape = (10, 10)
 
-    coords = (np.array([2, 4]), np.array([1,3]))
+    coords = (np.array([2, 4]), np.array([1, 3]))
 
     sampler = n2v_utils.pm_normal_withoutCP(1)
 
@@ -68,7 +70,7 @@ def test_pm_normal_neighbor_withoutCP():
     patch = np.arange(1000)
     patch.shape = (10, 10, 10, 1)
 
-    coords = (np.array([2, 4, 6]), np.array([1,3,5]), np.array([3,5,1]))
+    coords = (np.array([2, 4, 6]), np.array([1, 3, 5]), np.array([3, 5, 1]))
 
     for i in range(100):
         val = sampler(patch, coords, len(patch.shape))
@@ -119,7 +121,7 @@ def test_pm_normal_additive():
 
     val = sampler(patch, coords, len(patch.shape))
     for v, z, y, x in zip(val, *coords):
-        assert v == patch[z,y,x]
+        assert v == patch[z, y, x]
 
 
 def test_pm_normal_fitted():
@@ -162,4 +164,12 @@ def test_pm_identity():
 
     val = sampler(patch, coords, len(patch.shape))
     for v, z, y, x in zip(val, *coords):
-        assert v == patch[z, y, x]
+        assert v == patch[z, y, x]
+
+
+def test_tta():
+    img, _ = np.meshgrid(range(200), range(100))
+    img[:50, :50] = 50
+    aug = tta_forward(img[..., np.newaxis])
+    avg = tta_backward(aug)
+    assert np.sum(avg[..., 0] - img) == 0