diff --git a/competition/utils.py b/competition/utils.py old mode 100644 new mode 100755 index 1b19a36..19ad083 --- a/competition/utils.py +++ b/competition/utils.py @@ -41,25 +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 @@ -71,38 +101,58 @@ def nonlinear_transformation(x, prob=0.5): 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) @@ -111,8 +161,7 @@ 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, @@ -122,19 +171,27 @@ def image_in_painting(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 @@ -145,18 +202,56 @@ 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])] @@ -164,25 +259,7 @@ def generate_pair(img, batch_size, config, status="test"): 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) \ No newline at end of file + # apply augmentations + x[n], y[n] = generate_single_pair(y[n], config=config) + + yield (x, y)