def make_slices_iter(ids,
load_x,
load_y,
batch_size,
*,
shuffle=False,
empty_slice=True):
if shuffle:
ids = shuffle_ids(ids)
def slicer(ids):
for id_ in ids:
x, y = load_x(id_), load_y(id_)
yield from iterate_slices(x, y, empty=empty_slice)
# @pdp.pack_args
# def flip(mscan, segm):
# mscan = mscan.copy()
# segm = segm.copy()
# if np.random.randint(0, 2):
# mscan = np.flip(mscan, 1)
# segm = np.flip(segm, 1)
# if np.random.randint(0, 2):
# mscan = np.flip(mscan, 2)
# segm = np.flip(segm, 2)
# # randomly swap channels because why not?
# mscan = mscan[np.random.permutation(np.arange(len(mscan)))]
# return mscan, segm
return pdp.Pipeline(
pdp.Source(slicer(ids), buffer_size=30),
# pdp.LambdaTransformer(flip, buffer_size=3),
pdp.Chunker(chunk_size=batch_size, buffer_size=2),
pdp.LambdaTransformer(pdp.combine_batches, buffer_size=3))
def make_simple_iter(
ids, load_x, load_y, batch_size, *, shuffle=False):
if shuffle:
ids = shuffle_ids(ids)
return pdp.Pipeline(
pdp.Source(load_by_ids(load_x, load_y, ids), buffer_size=30),
pdp.Chunker(chunk_size=batch_size, buffer_size=2),
pdp.LambdaTransformer(pdp.combine_batches, buffer_size=3)
)
def make_3d_augm_patch_stratified_iter(
ids, load_x, load_y, *, batch_size,
x_patch_sizes, y_patch_size, nonzero_fraction, buffer_size=10):
x_patch_sizes = [np.array(x_patch_size) for x_patch_size in x_patch_sizes]
y_patch_size = np.array(y_patch_size)
spatial_dims = [-3, -2, -1]
random_seq = iter(partial(choice, ids), None)
def load_patient(name):
return Patient(name, load_x(name), load_y(name))
@lru_cache(maxsize=len(ids))
def find_cancer(patient: Patient):
if patient.y.ndim == 4:
mask = np.any(patient.y, axis=0)
elif patient.y.ndim == 3:
mask = patient.y > 0
else:
raise ValueError('segm number of dimensions '
'was {}, expected 3 or 4'.format(patient.y.ndim))
conditional_centre_indices = medim.patch.get_conditional_center_indices(
mask, patch_size=y_patch_size, spatial_dims=spatial_dims)
return patient.x, patient.y, conditional_centre_indices
@pdp.pack_args
def extract_big_patch(x, y, conditional_center_indices):
cancer_type = np.random.choice(
[True, False], p=[nonzero_fraction, 1 - nonzero_fraction])
if cancer_type:
center_idx = choice(conditional_center_indices)
else:
center_idx = medim.patch.get_uniform_center_index(
x_shape=np.array(x.shape), patch_size=y_patch_size,
spatial_dims=spatial_dims)
xs = medim.patch.extract_patch(
x, center_idx=center_idx, spatial_dims=spatial_dims,
patch_size=x_patch_sizes[0])
y = medim.patch.extract_patch(
y, center_idx=center_idx, patch_size=y_patch_size,
spatial_dims=spatial_dims)
return (xs, y)
def _scale_crop(x):
shape = np.array(x.shape)
x = zoom(x, scale, order=0)
new_shape = np.array(x.shape)
delta = shape - new_shape
d_pad = np.maximum(0, delta)
d_pad = list(zip(d_pad // 2, (d_pad + 1) // 2))
d_slice = np.maximum(0, -delta)
d_slice = zip(d_slice // 2, new_shape - (d_slice + 1) // 2)
d_slice = [slice(x, y) for x, y in d_slice]
x = x[d_slice]
x = np.pad(x, d_pad, mode='constant')
return x
def _rotate(x, order=3, theta=0, alpha=0):
x = rotate(x, theta, axes=(len(x.shape) - 2, len(x.shape) - 3),
reshape=False, order=order)
x = rotate(x, alpha, axes=(len(x.shape) - 2, len(x.shape) - 1),
reshape=False, order=order)
return x
def augment(x: np.ndarray, y: np.ndarray):
"""
Data random augmentation include scaling, rotation,
axes flipping.
__________________
Say thanks to Max!
"""
scipy.random.seed()
# scale = np.random.normal(1, 0.1, size=3)
alpha, theta = np.random.normal(0, 9, size=2)
alpha = 0
for i in range(1, len(x.shape) - 1):
if np.random.binomial(1, .5):
x = np.flip(x, -i)
y = np.flip(y, -i)
# mscan = np.array([_scale_crop(i) for i in mscan])
# segm = _scale_crop(segm[0])[np.newaxis]
x = _rotate(x, 3, theta, alpha)
y = _rotate(y, 0, theta, alpha)
# if np.random.binomial(1, .5):
# t = np.random.choice([-90, 0, 90])
# a = np.random.choice([-90, 0, 90])
# mscan = _rotate(mscan, 3, t, a)
# segm = _rotate(segm, 3, t, a)
x = np.array([i * np.random.normal(1, 0.35) for i in x])
return x, y
@pdp.pack_args
def augmentation(x_big, y):
return augment(x_big, y)
@pdp.pack_args
def extract_patch(x_big, y):
center_idx = np.array(x_big.shape)[spatial_dims] // 2 + np.array(
x_big.shape)[spatial_dims] % 2
xs = [x_big] + [medim.patch.extract_patch(
x_big, center_idx=center_idx, spatial_dims=spatial_dims,
patch_size=patch_size)
for patch_size in x_patch_sizes[1:]]
# segm = medim.patch.extract_patch(
# segm, center_idx=center_idx, patch_size=y_patch_size,
# spatial_dims=spatial_dims)
return (*xs, y)
return pdp.Pipeline(
pdp.Source(random_seq, buffer_size=3),
pdp.LambdaTransformer(load_patient, buffer_size=3),
pdp.LambdaTransformer(find_cancer, n_workers=8, buffer_size=3),
pdp.LambdaTransformer(extract_big_patch,
n_workers=8, buffer_size=batch_size),
pdp.LambdaTransformer(augmentation, n_workers=8,
buffer_size=0),
pdp.LambdaTransformer(extract_patch, n_workers=8,
buffer_size=batch_size),
pdp.Chunker(chunk_size=batch_size, buffer_size=3),
pdp.LambdaTransformer(pdp.combine_batches, buffer_size=buffer_size)
)
def make_3d_patch_stratified_iter(ids,
load_x,
load_y,
*,
batch_size,
x_patch_sizes,
y_patch_size,
nonzero_fraction,
buffer_size=10):
x_patch_sizes = [np.array(x_patch_size) for x_patch_size in x_patch_sizes]
y_patch_size = np.array(y_patch_size)
spatial_dims = [-3, -2, -1]
random_seq = iter(partial(choice, ids), None)
@lru_cache(maxsize=len(ids))
def load_patient(name):
return Patient(name, load_x(name), load_y(name))
@lru_cache(maxsize=len(ids))
def find_cancer(patient: Patient):
if len(patient.y.shape) == 3:
mask = patient.y > 0
elif len(patient.y.shape) == 4:
mask = np.any(patient.y, axis=0)
else:
raise ValueError('wrong number of dimensions ')
conditional_centre_indices = medim.patch.get_conditional_center_indices(
mask, patch_size=y_patch_size, spatial_dims=spatial_dims)
return patient.x, patient.y, conditional_centre_indices
@pdp.pack_args
def extract_patch(x, y, conditional_center_indices):
assert all([x.shape[i] == y.shape[i] for i in spatial_dims])
if np.random.uniform() < nonzero_fraction:
center_idx = choice(conditional_center_indices)
else:
center_idx = medim.patch.get_uniform_center_index(
x_shape=np.array(x.shape),
patch_size=y_patch_size,
spatial_dims=spatial_dims)
xs = [
medim.patch.extract_patch(x,
center_idx=center_idx,
spatial_dims=spatial_dims,
patch_size=patch_size)
for patch_size in x_patch_sizes
]
y = medim.patch.extract_patch(y,
center_idx=center_idx,
patch_size=y_patch_size,
spatial_dims=spatial_dims)
return (*xs, y)
return pdp.Pipeline(
pdp.Source(random_seq, buffer_size=3),
pdp.LambdaTransformer(load_patient, buffer_size=100),
pdp.LambdaTransformer(find_cancer, buffer_size=100),
pdp.LambdaTransformer(extract_patch,
n_workers=4,
buffer_size=batch_size),
pdp.Chunker(chunk_size=batch_size, buffer_size=3),
pdp.LambdaTransformer(pdp.combine_batches, buffer_size=buffer_size))