Skip to content

Commit

Permalink
Merge pull request #49 from joeranbosma/master
Browse files Browse the repository at this point in the history 
Major refactoring and bugfixes to allow data preprocessed by nnUNet
  • Loading branch information
@MrGiovanni
MrGiovanni authored Feb 18, 2022
2 parents e3a5ae3 + 24ae94c commit ece0390
Showing 1 changed file with 145 additions and 67 deletions.
212 changes: 145 additions & 67 deletions competition/utils.py
100644 → 100755
View file
Original file line number Diff line number Diff line change
@@ -1,25 +1,21 @@
from __future__ import print_function
import math
import os
import random
import copy
import scipy
import imageio
import string
import numpy as np
from skimage.transform import resize
try: # SciPy >= 0.19
from scipy.special import comb
except ImportError:
from scipy.misc import comb


def bernstein_poly(i, n, t):
"""
The Bernstein polynomial of n, i as a function of t
"""

return comb(n, i) * ( t**(n-i) ) * (1 - t)**i


def bezier_curve(points, nTimes=1000):
"""
Given a set of control points, return the
Expand All @@ -45,23 +41,55 @@ def bezier_curve(points, nTimes=1000):

return xvals, yvals

def data_augmentation(x, y, prob=0.5):

def random_flip_all_axes(x, y, prob=0.5):
"""
Randomly flip paired 3D images (e.g., scans and labels) along one or multiple axes
Input:
- x: 3D image, optionally with an additional axis for multiple modalities. Shape: (optional channel dimension), width, height, depth
- y: image to transform in the same way as x. Shape: same as x.
- prob: flip probability for each of the three iterations
"""
# augmentation by flipping
cnt = 3
while random.random() < prob and cnt > 0:
degree = random.choice([0, 1, 2])
degree = random.choice([-1, -2, -3])
x = np.flip(x, axis=degree)
y = np.flip(y, axis=degree)
cnt = cnt - 1

return x, y

def nonlinear_transformation(x, prob=0.5):

def nonlinear_transformation(x, prob=0.5, normalisation="z-score"):
"""
Perform nonlinear intensity transformation to input image
Input:
- x: 3D image, optionally with an additional axis for multiple modalities. Shape: (optional channel dimension), width, height, depth
- prob: transformation probability
Returns:
- nonlinear_x: (with probability prob) intensity transformed image of same shape as x
(with probability 1-prob) x
"""
if random.random() >= prob:
return x
points = [[0, 0], [random.random(), random.random()], [random.random(), random.random()], [1, 1]]
xpoints = [p[0] for p in points]
ypoints = [p[1] for p in points]

if normalisation == "z-score":
x_start, x_end = np.percentile(x, 0.1), np.percentile(x, 99.9)
points = [
[x_start, x_start],
[np.random.uniform(x_start, x_end), np.random.uniform(x_start, x_end)],
[np.random.uniform(x_start, x_end), np.random.uniform(x_start, x_end)],
[x_end, x_end]
]
elif normalisation == "minmax":
points = [[0, 0], [random.random(), random.random()], [random.random(), random.random()], [1, 1]]
else:
raise ValueError(f"Unrecognised normalisation method: {normalisation}")

xvals, yvals = bezier_curve(points, nTimes=100000)
if random.random() < 0.5:
# Half change to get flip
Expand All @@ -71,38 +99,60 @@ def nonlinear_transformation(x, prob=0.5):
nonlinear_x = np.interp(x, xvals, yvals)
return nonlinear_x


def local_pixel_shuffling(x, prob=0.5):
if random.random() >= prob:
return x
"""
Perform local pixel shuffling to input image, with probability prob
Input:
- x: 3D image, with an additional axis for multiple modalities. Shape: channel dimension, width, height, depth
- prob: transformation probability per modality
Returns:
- local_shuffling_x: transformed image of same shape as x
"""
image_temp = copy.deepcopy(x)
orig_image = copy.deepcopy(x)
_, img_rows, img_cols, img_deps = x.shape
channels, img_rows, img_cols, img_deps = x.shape
num_block = 10000
for _ in range(num_block):
block_noise_size_x = random.randint(1, img_rows//10)
block_noise_size_y = random.randint(1, img_cols//10)
block_noise_size_z = random.randint(1, img_deps//10)
noise_x = random.randint(0, img_rows-block_noise_size_x)
noise_y = random.randint(0, img_cols-block_noise_size_y)
noise_z = random.randint(0, img_deps-block_noise_size_z)
window = orig_image[0, noise_x:noise_x+block_noise_size_x,
noise_y:noise_y+block_noise_size_y,
noise_z:noise_z+block_noise_size_z,
]
window = window.flatten()
np.random.shuffle(window)
window = window.reshape((block_noise_size_x,
block_noise_size_y,
block_noise_size_z))
image_temp[0, noise_x:noise_x+block_noise_size_x,
noise_y:noise_y+block_noise_size_y,
noise_z:noise_z+block_noise_size_z] = window
for i in range(channels):
if random.random() >= prob:
continue
for _ in range(num_block):
block_noise_size_x = random.randint(1, img_rows//10)
block_noise_size_y = random.randint(1, img_cols//10)
block_noise_size_z = random.randint(1, img_deps//10)
noise_x = random.randint(0, img_rows-block_noise_size_x)
noise_y = random.randint(0, img_cols-block_noise_size_y)
noise_z = random.randint(0, img_deps-block_noise_size_z)
window = orig_image[0, noise_x:noise_x+block_noise_size_x,
noise_y:noise_y+block_noise_size_y,
noise_z:noise_z+block_noise_size_z,
]
window = window.flatten()
np.random.shuffle(window)
window = window.reshape((block_noise_size_x,
block_noise_size_y,
block_noise_size_z))
image_temp[i, noise_x:noise_x+block_noise_size_x,
noise_y:noise_y+block_noise_size_y,
noise_z:noise_z+block_noise_size_z] = window
local_shuffling_x = image_temp

return local_shuffling_x


def image_in_painting(x):
_, img_rows, img_cols, img_deps = x.shape
"""
Perform image inpainting to input image
Input:
- x: 3D image. Shape: width, height, depth
Returns:
- x: transformed image of same shape as x
"""
img_rows, img_cols, img_deps = x.shape
cnt = 5
while cnt > 0 and random.random() < 0.95:
block_noise_size_x = random.randint(img_rows//6, img_rows//3)
Expand All @@ -111,29 +161,37 @@ def image_in_painting(x):
noise_x = random.randint(3, img_rows-block_noise_size_x-3)
noise_y = random.randint(3, img_cols-block_noise_size_y-3)
noise_z = random.randint(3, img_deps-block_noise_size_z-3)
x[:,
noise_x:noise_x+block_noise_size_x,
x[noise_x:noise_x+block_noise_size_x,
noise_y:noise_y+block_noise_size_y,
noise_z:noise_z+block_noise_size_z] = np.random.rand(block_noise_size_x,
block_noise_size_y,
block_noise_size_z, ) * 1.0
cnt -= 1
return x


def image_out_painting(x):
_, img_rows, img_cols, img_deps = x.shape
"""
Perform image outpainting to input image
Input:
- x: 3D image. Shape: width, height, depth
Returns:
- x: transformed image of same shape as x
"""
img_rows, img_cols, img_deps = x.shape
image_temp = copy.deepcopy(x)
x = np.random.rand(x.shape[0], x.shape[1], x.shape[2], x.shape[3], ) * 1.0
x = np.random.rand(*x.shape) * 1.0
block_noise_size_x = img_rows - random.randint(3*img_rows//7, 4*img_rows//7)
block_noise_size_y = img_cols - random.randint(3*img_cols//7, 4*img_cols//7)
block_noise_size_z = img_deps - random.randint(3*img_deps//7, 4*img_deps//7)
noise_x = random.randint(3, img_rows-block_noise_size_x-3)
noise_y = random.randint(3, img_cols-block_noise_size_y-3)
noise_z = random.randint(3, img_deps-block_noise_size_z-3)
x[:,
noise_x:noise_x+block_noise_size_x,
x[noise_x:noise_x+block_noise_size_x,
noise_y:noise_y+block_noise_size_y,
noise_z:noise_z+block_noise_size_z] = image_temp[:, noise_x:noise_x+block_noise_size_x,
noise_z:noise_z+block_noise_size_z] = image_temp[noise_x:noise_x+block_noise_size_x,
noise_y:noise_y+block_noise_size_y,
noise_z:noise_z+block_noise_size_z]
cnt = 4
Expand All @@ -144,44 +202,64 @@ def image_out_painting(x):
noise_x = random.randint(3, img_rows-block_noise_size_x-3)
noise_y = random.randint(3, img_cols-block_noise_size_y-3)
noise_z = random.randint(3, img_deps-block_noise_size_z-3)
x[:,
noise_x:noise_x+block_noise_size_x,
x[noise_x:noise_x+block_noise_size_x,
noise_y:noise_y+block_noise_size_y,
noise_z:noise_z+block_noise_size_z] = image_temp[:, noise_x:noise_x+block_noise_size_x,
noise_z:noise_z+block_noise_size_z] = image_temp[noise_x:noise_x+block_noise_size_x,
noise_y:noise_y+block_noise_size_y,
noise_z:noise_z+block_noise_size_z]
cnt -= 1
return x


def generate_single_pair(y, config):
"""
Generate data augmentations to a single image with probabilities specified in the config
Input:
- y: original 3D image, with an additional axis for multiple modalities. Shape: channel dimension, width, height, depth.
Returns:
- x: transformed 3D image. Shape: channel dimension, width, height, depth.
- y: original 3D image. Shape: channel dimension, width, height, depth.
"""
# Autoencoder
x = copy.deepcopy(y)

# Flip
x, y = random_flip_all_axes(x, y, config.flip_rate)

# Local Shuffle Pixel
x = local_pixel_shuffling(x, prob=config.local_rate)

# Apply non-Linear transformation with an assigned probability
x = nonlinear_transformation(x, config.nonlinear_rate)

# Inpainting & Outpainting
channels = x.shape[0]
for i in range(channels):
if random.random() < config.paint_rate:
if random.random() < config.inpaint_rate:
# Inpainting
x[i] = image_in_painting(x[i])
else:
# Outpainting
x[i] = image_out_painting(x[i])

return x, y


def generate_pair(img, batch_size, config, status="test"):
"""
img: Images, shape (bs; channel; z; y; x)
"""
img_rows, img_cols, img_deps = img.shape[2], img.shape[3], img.shape[4]
while True:
index = [i for i in range(img.shape[0])]
random.shuffle(index)
y = img[index[:batch_size]]
x = copy.deepcopy(y)
for n in range(batch_size):

# Autoencoder
x[n] = copy.deepcopy(y[n])

# Flip
x[n], y[n] = data_augmentation(x[n], y[n], config.flip_rate)

# Local Shuffle Pixel
x[n] = local_pixel_shuffling(x[n], prob=config.local_rate)

# Apply non-Linear transformation with an assigned probability
x[n] = nonlinear_transformation(x[n], config.nonlinear_rate)

# Inpainting & Outpainting
if random.random() < config.paint_rate:
if random.random() < config.inpaint_rate:
# Inpainting
x[n] = image_in_painting(x[n])
else:
# Outpainting
x[n] = image_out_painting(x[n])
yield (x, y)
# apply augmentations
x[n], y[n] = generate_single_pair(y[n], config=config)

yield (x, y)

0 comments on commit ece0390

Please sign in to comment.