def pad_nd_image_and_seg(data, seg, new_shape=None, must_be_divisible_by=None, pad_mode_data='constant',
np_pad_kwargs_data=None, pad_mode_seg='constant', np_pad_kwargs_seg=None):
"""
Pads data and seg to new_shape. new_shape is thereby understood as min_shape (if data/seg is already larger then
new_shape the shape stays the same for the dimensions this applies)
:param data:
:param seg:
:param new_shape: if none then only must_be_divisible_by is applied
:param must_be_divisible_by: UNet like architectures sometimes require the input to be divisibly by some number. This
will modify new_shape if new_shape is not divisibly by this (by increasing it accordingly).
must_be_divisible_by should be a list of int (one for each spatial dimension) and this list must have the same
length as new_shape
:param pad_mode_data: see np.pad
:param np_pad_kwargs_data:see np.pad
:param pad_mode_seg:see np.pad
:param np_pad_kwargs_seg:see np.pad
:return:
"""
assert len(new_shape) == len(data.shape), "data_shape and new_shape must have the same dimensionality"
sample_data = pad_nd_image(data, new_shape, mode=pad_mode_data, kwargs=np_pad_kwargs_data,
return_slicer=False, shape_must_be_divisible_by=must_be_divisible_by)
if seg is not None:
sample_seg = pad_nd_image(seg, new_shape, mode=pad_mode_seg, kwargs=np_pad_kwargs_seg,
return_slicer=False, shape_must_be_divisible_by=must_be_divisible_by)
else:
sample_seg = None
return sample_data, sample_seg
def preprocessing(self, sample, training=True):
# Access data
img_data = sample.img_data
seg_data = sample.seg_data
# Transform data from channel-last to channel-first structure
img_data = np.moveaxis(img_data, -1, 0)
if training: seg_data = np.moveaxis(seg_data, -1, 0)
# Pad imaging data
img_data, crop_coords = pad_nd_image(
img_data,
self.min_size,
mode=self.pad_mode,
kwargs={"constant_values": self.pad_value_img},
return_slicer=True,
shape_must_be_divisible_by=self.shape_must_be_divisible_by)
if training:
seg_data = pad_nd_image(
seg_data,
self.min_size,
mode=self.pad_mode,
kwargs={"constant_values": self.pad_value_seg},
return_slicer=False,
shape_must_be_divisible_by=self.shape_must_be_divisible_by)
# Cache current crop coordinates for later postprocessing
if not training: self.original_coords = crop_coords
# Transform data from channel-first back to channel-last structure
img_data = np.moveaxis(img_data, 0, -1)
if training: seg_data = np.moveaxis(seg_data, 0, -1)
# Save resampled imaging data to sample
sample.img_data = img_data
sample.seg_data = seg_data
def generate_train_batch(self):
subjects = self._data[0]
subject_idx = int(random.uniform(0, len(subjects)))
data, seg = load_training_data(self.Config, subjects[subject_idx])
# Convert peaks to tensors if tensor model
if self.Config.NR_OF_GRADIENTS == 18*self.Config.NR_SLICES:
data = peak_utils.peaks_to_tensors(data)
slice_direction = data_utils.slice_dir_to_int(self.Config.TRAINING_SLICE_DIRECTION)
if data.shape[slice_direction] <= self.batch_size:
print("INFO: Batch size bigger than nr of slices. Therefore sampling with replacement.")
slice_idxs = np.random.choice(data.shape[slice_direction], self.batch_size, True, None)
else:
slice_idxs = np.random.choice(data.shape[slice_direction], self.batch_size, False, None)
if self.Config.NR_SLICES > 1:
x, y = data_utils.sample_Xslices(data, seg, slice_idxs, slice_direction=slice_direction,
labels_type=self.Config.LABELS_TYPE, slice_window=self.Config.NR_SLICES)
else:
x, y = data_utils.sample_slices(data, seg, slice_idxs, slice_direction=slice_direction,
labels_type=self.Config.LABELS_TYPE)
# Can be replaced by crop
# x = pad_nd_image(x, self.Config.INPUT_DIM, mode='constant', kwargs={'constant_values': 0})
# y = pad_nd_image(y, self.Config.INPUT_DIM, mode='constant', kwargs={'constant_values': 0})
# x = center_crop_2D_image_batched(x, self.Config.INPUT_DIM)
# y = center_crop_2D_image_batched(y, self.Config.INPUT_DIM)
# If want to convert e.g. 1.25mm (HCP) image to 2mm image (bb)
# x, y = self._zoom_x_and_y(x, y, 0.67) # very slow -> try spatial_transform, should be fast
if self.Config.PAD_TO_SQUARE:
#Crop and pad to input size
x, y = crop(x, y, crop_size=self.Config.INPUT_DIM) # does not work with img with batches and channels
else:
# Works -> results as good?
# Will pad each axis to be multiple of 16. (Each sample can end up having different dimensions. Also x and y
# can be different)
# This is needed for Schizo dataset
x = pad_nd_image(x, shape_must_be_divisible_by=(16, 16), mode='constant', kwargs={'constant_values': 0})
y = pad_nd_image(y, shape_must_be_divisible_by=(16, 16), mode='constant', kwargs={'constant_values': 0})
# Does not make it slower
x = x.astype(np.float32)
y = y.astype(np.float32)
# possible optimization: sample slices from different patients and pad all to same size (size of biggest)
data_dict = {"data": x, # (batch_size, channels, x, y, [z])
"seg": y,
"slice_dir": slice_direction} # (batch_size, channels, x, y, [z])
return data_dict
def __getitem__(self, item):
data_sample, seg_sample = self.get_training_data(self.cases_lists[item])
if self.is_training:
data_sample, seg_sample = do_augmentation_train(data_sample, seg_sample, patch_size=self.final_patch_size,
params=self.data_aug_params)
else:
data_sample, seg_sample = do_augmentation_test(data_sample, seg_sample, self.data_aug_params)
if self.final_patch_size is not None and any([i<j for i, j in zip(data_sample.shape[1:], self.final_patch_size)]):
data_sample = pad_nd_image(data_sample, self.final_patch_size, mode='constant', **{'constant_values': 0})
seg_sample = pad_nd_image(seg_sample, self.final_patch_size, mode='constant', **{'constant_values': 0})
return self.cases_lists[item], data_sample, seg_sample[0] # data (c, d, h, w), seg (d, h, w)
def generate_train_batch(self):
subjects = self._data[0]
# subject_idx = int(random.uniform(0, len(subjects))) # len(subjects)-1 not needed because int always rounds to floor
subject_idxs = np.random.choice(len(subjects), self.batch_size, False,
None)
x = []
y = []
for subject_idx in subject_idxs:
data, seg = load_training_data(
self.Config, subjects[subject_idx]) # (x, y, z, channels)
data = data.transpose(3, 0, 1, 2) # channels have to be first
seg = seg.transpose(3, 0, 1, 2)
x.append(data)
y.append(seg)
x = np.array(x)
y = np.array(y)
# Can be replaced by crop -> shorter
# x = pad_nd_image(x, self.Config.INPUT_DIM, mode='constant', kwargs={'constant_values': 0})
# y = pad_nd_image(y, self.Config.INPUT_DIM, mode='constant', kwargs={'constant_values': 0})
# x = center_crop_3D_image_batched(x, self.Config.INPUT_DIM)
# y = center_crop_3D_image_batched(y, self.Config.INPUT_DIM)
# Crop and pad to input size
# x, y = crop(x, y, crop_size=self.Config.INPUT_DIM) # does not work with img with batches and channels
# Works
# This is needed for Schizo dataset
x = pad_nd_image(x,
shape_must_be_divisible_by=(8, 8),
mode='constant',
kwargs={'constant_values': 0})
y = pad_nd_image(y,
shape_must_be_divisible_by=(8, 8),
mode='constant',
kwargs={'constant_values': 0})
x = x.astype(np.float32)
y = y.astype(np.float32)
data_dict = {
"data": x, # (batch_size, channels, x, y, [z])
"seg": y
} # (batch_size, channels, x, y, [z])
return data_dict
def predict_patient_in_patches(patient_data, model):
# we pad the patient data in order to fit the patches in it
patient_data_pd = pad_nd_image(patient_data, [144, 192, 192]) # 24*6, 128+2*32, 128+2*32
# patches.shape = (1, 1, 6, 3, 3, 1, 3, 24, 128, 128)
steps = (1,1,args.patch_depth,int(args.patch_width/4),int(args.patch_height/4))
window_shape = (1, 3, args.patch_depth, args.patch_width, args.patch_height)
patches = skimage.util.view_as_windows(patient_data_pd[:, :3, :, :, :], window_shape=window_shape, step=steps)
# (1, 4, 138, 169, 141)
target_shape = list(patient_data_pd.shape)
print(f"Target shape in predict patient in patches is: {target_shape}")
if args.multi_class:
target_shape[1] = 4
else:
target_shape[1] = 1 # only one output channel
prediction = torch.zeros(*target_shape)
if args.use_gpu:
prediction = prediction.cuda()
for i in range(patches.shape[2]):
for j in range(patches.shape[3]):
for k in range(patches.shape[4]):
data = torch.from_numpy(patches[0, 0, i, j, k])
if args.use_gpu:
data = data.cuda()
output = model.forward(data)
prediction[:, :,
i*steps[2]:i*steps[2]+window_shape[2],
j*steps[3]:j*steps[3]+window_shape[3],
k*steps[4]:k*steps[4]+window_shape[4]] += output
return prediction
def _internal_predict_3D_3Dconv(self, x: np.ndarray, min_size: Tuple[int, ...], do_mirroring: bool,
mirror_axes: tuple = (0, 1, 2), regions_class_order: tuple = None,
pad_border_mode: str = "constant", pad_kwargs: dict = None,
verbose: bool = True) -> Tuple[np.ndarray, np.ndarray]:
"""
This one does fully convolutional inference. No sliding window
"""
assert len(x.shape) == 4, "x must be (c, x, y, z)"
assert self.get_device() != "cpu"
assert self.input_shape_must_be_divisible_by is not None, 'input_shape_must_be_divisible_by must be set to ' \
'run _internal_predict_3D_3Dconv'
if verbose: print("do mirror:", do_mirroring)
data, slicer = pad_nd_image(x, min_size, pad_border_mode, pad_kwargs, True,
self.input_shape_must_be_divisible_by)
predicted_probabilities = self._internal_maybe_mirror_and_pred_3D(data[None], mirror_axes, do_mirroring,
None)[0]
slicer = tuple(
[slice(0, predicted_probabilities.shape[i]) for i in range(len(predicted_probabilities.shape) -
(len(slicer) - 1))] + slicer[1:])
predicted_probabilities = predicted_probabilities[slicer]
if regions_class_order is None:
predicted_segmentation = predicted_probabilities.argmax(0)
predicted_segmentation = predicted_segmentation.detach().cpu().numpy()
predicted_probabilities = predicted_probabilities.detach().cpu().numpy()
else:
predicted_probabilities = predicted_probabilities.detach().cpu().numpy()
predicted_segmentation = np.zeros(predicted_probabilities.shape[1:], dtype=np.float32)
for i, c in enumerate(regions_class_order):
predicted_segmentation[predicted_probabilities[i] > 0.5] = c
return predicted_segmentation, predicted_probabilities
def generate_train_batch(self):
# DataLoader has its own methods for selecting what patients to use next, see its Documentation
idx = self.get_indices()
patients_for_batch = [self._data[i] for i in idx]
# initialize empty array for data and seg
data = np.zeros((self.batch_size, self.num_modalities,
*self.patch_size), dtype=np.float32)
seg = np.zeros((self.batch_size, 1, *self.patch_size),
dtype=np.float32)
metadata = []
patient_names = []
# iterate over patients_for_batch and include them in the batch
for i, j in enumerate(patients_for_batch):
patient_data, patient_metadata = self.load_patient(j)
# this will only pad patient_data if its shape is smaller than self.patch_size
patient_data = pad_nd_image(patient_data, self.patch_size)
# now random crop to self.patch_size
# crop expects the data to be (b, c, x, y, z) but patient_data is (c, x, y, z) so we need to add one
# dummy dimension in order for it to work (@Todo, could be improved)
patient_data, patient_seg = crop(
patient_data[:-1][None], patient_data[-1:][None], self.patch_size, crop_type="random")
data[i] = patient_data[0]
seg[i] = patient_seg[0]
metadata.append(patient_metadata)
patient_names.append(j)
return {'data': data, 'seg': seg, 'metadata': metadata, 'names': patient_names}
def _internal_predict_3D_3Dconv(self, x, do_mirroring, num_repeats, min_size=None, BATCH_SIZE=None,
mirror_axes=(0, 1, 2), regions_class_order=None, pad_border_mode="edge",
pad_kwargs=None):
with torch.no_grad():
x, slicer = pad_nd_image(x, min_size, pad_border_mode, pad_kwargs, True, self.input_shape_must_be_divisible_by)
#x, old_shape = pad_patient_3D_incl_c(x, self.input_shape_must_be_divisible_by, min_size)
new_shp = x.shape
data = np.zeros(tuple([1] + list(new_shp)), dtype=np.float32)
data[0] = x
if BATCH_SIZE is not None:
data = np.vstack([data] * BATCH_SIZE)
stacked = self._internal_maybe_mirror_and_pred_3D(data, num_repeats, mirror_axes, do_mirroring, None)[0]
slicer = tuple([slice(0, stacked.shape[i]) for i in range(len(stacked.shape) - (len(slicer) - 1))] + slicer[1:])
stacked = stacked[slicer]
softmax_pred = stacked
if regions_class_order is None:
predicted_segmentation = softmax_pred.argmax(0)
else:
predicted_segmentation_shp = softmax_pred[0].shape
predicted_segmentation = np.zeros(predicted_segmentation_shp, dtype=np.float32)
for i, c in enumerate(regions_class_order):
predicted_segmentation[softmax_pred[i] > 0.5] = c
return predicted_segmentation, None, softmax_pred, None
def preprocess_MLPerf(model, checkpoint_name, folds, fp16, list_of_lists, output_filenames, preprocessing_folder, num_threads_preprocessing):
assert len(list_of_lists) == len(output_filenames)
print("loading parameters for folds", folds)
trainer, params = load_model_and_checkpoint_files(model, folds, fp16=fp16, checkpoint_name=checkpoint_name)
print("starting preprocessing generator")
preprocessing = preprocess_multithreaded(trainer, list_of_lists, output_filenames, num_threads_preprocessing, None)
print("Preprocessing images...")
all_output_files = []
for preprocessed in preprocessing:
output_filename, (d, dct) = preprocessed
all_output_files.append(output_filename)
if isinstance(d, str):
data = np.load(d)
os.remove(d)
d = data
# Pad to the desired full volume
d = pad_nd_image(d, trainer.patch_size, "constant", None, False, None)
with open(os.path.join(preprocessing_folder, output_filename+ ".pkl"), "wb") as f:
pickle.dump([d, dct], f)
f.close()
return all_output_files
def preprocess_data(root_dir, y_shape=64, z_shape=64):
image_dir = os.path.join(root_dir, 'imagesTr')
label_dir = os.path.join(root_dir, 'labelsTr')
output_dir = os.path.join(root_dir, 'preprocessed')
classes = 3
if not os.path.exists(output_dir):
os.makedirs(output_dir)
print('Created' + output_dir + '...')
class_stats = defaultdict(int)
total = 0
nii_files = subfiles(image_dir, suffix=".nii.gz", join=False)
for i in range(0, len(nii_files)):
if nii_files[i].startswith("._"):
nii_files[i] = nii_files[i][2:]
for f in nii_files:
image, _ = load(os.path.join(image_dir, f))
label, _ = load(os.path.join(label_dir, f.replace('_0000', '')))
print(f)
for i in range(classes):
class_stats[i] += np.sum(label == i)
total += np.sum(label == i)
# normalize images
image = (image - image.min()) / (image.max() - image.min())
image = pad_nd_image(image, (image.shape[0], y_shape, z_shape),
"constant",
kwargs={'constant_values': image.min()})
label = pad_nd_image(label, (image.shape[0], y_shape, z_shape),
"constant",
kwargs={'constant_values': label.min()})
result = np.stack((image, label))
np.save(os.path.join(output_dir, f.split('.')[0] + '.npy'), result)
print(f)
print(total)
for i in range(classes):
print(class_stats[i], class_stats[i] / total)
def generate_train_batch(self):
idx = self.get_indices(
) # get batch_size subjects randomly, return a list
data = np.zeros(
(self.batch_size, self.nb_modalities, *self.patch_size),
dtype=np.float32)
seg = np.zeros((self.batch_size, 1, *self.patch_size), dtype=np.int64)
for i, j in enumerate(idx):
tmp_data = pad_nd_image(self._data[j]['img'], self.patch_size)
tmp_seg = pad_nd_image(self._data[j]['seg'], self.patch_size)
cropped_data, cropped_seg = crop(tmp_data,
tmp_seg,
self.patch_size,
crop_type='random')
data[i] = cropped_data[0]
seg[i] = cropped_seg[0]
return {'data': data, 'seg': seg} # (b, 2, d, h, w), (b, 1, d, h, w)
def get_data_from_array(self, open_array):
data = []
fnames = []
slice_idxs = []
labels = []
for slice in open_array:
fn_name = self.files[slice[0]]
numpy_array = np.load(fn_name, mmap_mode="r")
numpy_slice = numpy_array[:, slice[1]]
print("numpy_array shape :", numpy_array.shape,
"numpy_slice shape :", numpy_slice.shape)
# this will only pad patient_data if its shape is smaller than self.patch_size
patient_data = pad_nd_image(numpy_slice, self.patch_size)
# now random crop to self.patch_size
# crop expects the data to be (b, c, x, y, z) but patient_data is (c, x, y, z) so we need to add one
# dummy dimension in order for it to work (@Todo, could be improved)
patient_data, patient_seg = crop(patient_data[:-1][None],
patient_data[-1:][None],
self.patch_size,
crop_type="random")
dataslice = patient_data[0][0][None]
dataslice = dataslice[0][None]
segGt = patient_seg[0]
#print("Before data shape :",dataslice.shape ,"sdpegGt :",segGt.shape)
#dataslice = dataslice[0]
print("After data shape :", dataslice.shape, "segGt :",
segGt.shape)
data.append(dataslice) # 'None' keeps the dimension
if self.label_slice is not None:
labels.append(segGt) # 'None' keeps the dimension
fnames.append(self.files[slice[0]])
slice_idxs.append(slice[1])
ret_dict = {'data': data, 'fnames': fnames, 'slice_idxs': slice_idxs}
if self.label_slice is not None:
ret_dict['seg'] = labels
print("data shape :", len(data), "labels :", len(labels))
return ret_dict
def pad_patch(patch, patch_shape, return_slicer=False):
# Initialize stat length to overwrite batchgenerators default
kwargs = {"stat_length": None}
# Transform prediction from channel-last to channel-first structure
patch = np.moveaxis(patch, -1, 1)
# Run padding
padding_results = pad_nd_image(patch, new_shape=patch_shape,
mode="minimum", return_slicer=return_slicer,
kwargs=kwargs)
# Return padding results
if return_slicer:
# Transform data from channel-first back to channel-last structure
padded_patch = np.moveaxis(padding_results[0], 1, -1)
return padded_patch, padding_results[1]
else:
# Transform data from channel-first back to channel-last structure
padding_results = np.moveaxis(padding_results, 1, -1)
return padding_results
def generate_train_batch(self):
idx = self.get_indices()
patients_for_batch = [self._data[i] for i in idx]
data = np.zeros((self.batch_size, self.num_modalities, *self.patch_size), dtype=np.float32)
for i, _ in enumerate(patients_for_batch):
patient_data = patients_for_batch[i]
if not self.crop:
patient_data = resize(patient_data, (patient_data.shape[0],)+self.patch_size,
anti_aliasing=True)
data[i] = patient_data
else:
# this will only pad patient_data if its shape is smaller than self.patch_size
patient_data = pad_nd_image(patient_data, self.patch_size)
# now random crop to self.patch_size
patient_data = crop(patient_data[None], crop_size=self.patch_size, crop_type="center")
data[i] = patient_data[0]
return {'data': data}
def _internal_predict_3D_3Dconv_tiled(self,
x,
num_repeats,
BATCH_SIZE=None,
tile_in_z=True,
step=2,
do_mirroring=True,
mirror_axes=(0, 1, 2),
patch_size=None,
regions_class_order=None,
use_gaussian=False,
pad_border_mode="edge",
pad_kwargs=None,
all_in_gpu=False,
nb_of_classes=None,
return_softmax=True):
"""
x must be (c, x, y, z)
:param x:
:param num_repeats:
:param BATCH_SIZE:
:param tile_in_z:
:param step:
:param do_mirroring:
:param mirror_axes:
:param patch_size:
:param regions_class_order:
:param use_gaussian:
:param pad_border_mode:
:param pad_kwargs:
:param all_in_gpu: if True then data and prediction will be held in GPU for inference. Faster, but uses more vram
:param return_softmax: if all_in_gpu, the copying of the softmax back to cpu can take a while. If you dont
need the softmax (just the segmentation) and all_in_gpu=True then you can shave a few seconds off the inference
time by setting this to False
:return:
"""
assert len(x.shape) == 4, "x must be (c, x, y, z)"
assert self.get_device() != "cpu"
assert not all_in_gpu, "all_in_gpu is currently broken. Rounding errors in fp16... needs to be reimplemented with fp32"
print("step:", step)
print("do mirror:", do_mirroring)
torch.cuda.empty_cache()
with torch.no_grad():
assert patch_size is not None, "patch_size cannot be None for tiled prediction"
data, slicer = pad_nd_image(x, patch_size, pad_border_mode,
pad_kwargs, True, None)
data = data[None]
if BATCH_SIZE is not None:
data = np.vstack([data] * BATCH_SIZE)
input_size = [1, x.shape[0]] + list(patch_size)
if not tile_in_z:
input_size[2] = data.shape[2]
patch_size[0] = data.shape[2]
input_size = [int(i) for i in input_size]
if nb_of_classes is None:
a = torch.zeros(input_size,
dtype=torch.float).cuda(self.get_device(),
non_blocking=True)
# dummy run to see number of classes
nb_of_classes = self(a).size()[1]
if use_gaussian:
if self._gaussian_3d is None or not all([
i == j for i, j in zip(
patch_size, self._patch_size_for_gaussian_3d)
]):
tmp = np.zeros(patch_size)
center_coords = [i // 2 for i in patch_size]
sigmas = [i / 8 for i in patch_size]
tmp[tuple(center_coords)] = 1
tmp_smooth = gaussian_filter(tmp,
sigmas,
0,
mode='constant',
cval=0)
tmp_smooth = tmp_smooth / tmp_smooth.max() * 1
add = tmp_smooth # + 1e-7
# tmp_smooth cannot be 0, otherwise we may end up with nans!
tmp_smooth[tmp_smooth == 0] = np.min(
tmp_smooth[tmp_smooth != 0])
# we only need to compute that once. It can take a while to compute this due to the large sigma in
# gaussian_filter
self._gaussian_3d = np.copy(add)
self._patch_size_for_gaussian_3d = deepcopy(patch_size)
else:
print("using precomputed Gaussian")
add = self._gaussian_3d
else:
add = np.ones(patch_size, dtype=np.float32)
add = add.astype(np.float32)
data_shape = data.shape
print("configuring tiles")
center_coord_start = np.array([i // 2
for i in patch_size]).astype(int)
center_coord_end = np.array([
data_shape[i + 2] - patch_size[i] // 2
for i in range(len(patch_size))
]).astype(int)
num_steps = np.ceil([
(center_coord_end[i] - center_coord_start[i]) /
(patch_size[i] / step) for i in range(3)
])
step_size = np.array([
(center_coord_end[i] - center_coord_start[i]) /
(num_steps[i] + 1e-8) for i in range(3)
])
step_size[step_size == 0] = 9999999
xsteps = np.round(
np.arange(center_coord_start[0], center_coord_end[0] + 1e-8,
step_size[0])).astype(int)
ysteps = np.round(
np.arange(center_coord_start[1], center_coord_end[1] + 1e-8,
step_size[1])).astype(int)
zsteps = np.round(
np.arange(center_coord_start[2], center_coord_end[2] + 1e-8,
step_size[2])).astype(int)
if all_in_gpu:
print("initializing result array (on GPU)")
# some of these can remain in half. We just need the reuslts for softmax so it won't hurt at all to reduce
# precision. Inference is of course done in float
result = torch.zeros([nb_of_classes] + list(data.shape[2:]),
dtype=torch.half,
device=self.get_device())
print("moving data to GPU")
data = torch.from_numpy(data).cuda(self.get_device(),
non_blocking=True)
print("initializing result_numsamples (on GPU)")
result_numsamples = torch.zeros([nb_of_classes] +
list(data.shape[2:]),
dtype=torch.half,
device=self.get_device())
print("moving add to GPU")
add = torch.from_numpy(add).cuda(self.get_device(),
non_blocking=True).half()
add_torch = add
else:
result = np.zeros([nb_of_classes] + list(data.shape[2:]),
dtype=np.float32)
result_numsamples = np.zeros([nb_of_classes] +
list(data.shape[2:]),
dtype=np.float32)
add_torch = torch.from_numpy(add).cuda(self.get_device(),
non_blocking=True)
print("data shape:", data_shape)
print("patch size:", patch_size)
print("steps (x, y, and z):", xsteps, ysteps, zsteps)
print("number of tiles:", len(xsteps) * len(ysteps) * len(zsteps))
# data, result and add_torch and result_numsamples are now on GPU
for x in xsteps:
lb_x = x - patch_size[0] // 2
ub_x = x + patch_size[0] // 2
for y in ysteps:
lb_y = y - patch_size[1] // 2
ub_y = y + patch_size[1] // 2
for z in zsteps:
lb_z = z - patch_size[2] // 2
ub_z = z + patch_size[2] // 2
predicted_patch = \
self._internal_maybe_mirror_and_pred_3D(data[:, :, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z],
num_repeats, mirror_axes, do_mirroring, add_torch)[0]
if all_in_gpu:
predicted_patch = predicted_patch.half()
else:
predicted_patch = predicted_patch.cpu().numpy()
result[:, lb_x:ub_x, lb_y:ub_y,
lb_z:ub_z] += predicted_patch
if all_in_gpu:
result_numsamples[:, lb_x:ub_x, lb_y:ub_y,
lb_z:ub_z] += add.half()
else:
result_numsamples[:, lb_x:ub_x, lb_y:ub_y,
lb_z:ub_z] += add
slicer = tuple([
slice(0, result.shape[i])
for i in range(len(result.shape) - (len(slicer) - 1))
] + slicer[1:])
result = result[slicer]
result_numsamples = result_numsamples[slicer]
softmax_pred = result / result_numsamples
# patient_data = patient_data[:, :old_shape[0], :old_shape[1], :old_shape[2]]
if regions_class_order is None:
predicted_segmentation = softmax_pred.argmax(0)
else:
if all_in_gpu:
softmax_pred_here = softmax_pred.detach().cpu().numpy()
else:
softmax_pred_here = softmax_pred
predicted_segmentation_shp = softmax_pred_here[0].shape
predicted_segmentation = np.zeros(predicted_segmentation_shp,
dtype=np.float32)
for i, c in enumerate(regions_class_order):
predicted_segmentation[softmax_pred_here[i] > 0.5] = c
if all_in_gpu:
print("copying results to CPU")
predicted_segmentation = predicted_segmentation.detach().cpu(
).numpy()
if return_softmax:
softmax_pred = softmax_pred.half().detach().cpu().numpy()
else:
softmax_pred = None
print("prediction on GPU done")
return predicted_segmentation, None, softmax_pred, None
def _internal_predict_2D_2Dconv_tiled(
self, x: np.ndarray, step_size: float, do_mirroring: bool,
mirror_axes: tuple, patch_size: tuple, regions_class_order: tuple,
use_gaussian: bool, pad_border_mode: str, pad_kwargs: dict,
all_in_gpu: bool, verbose: bool) -> Tuple[np.ndarray, np.ndarray]:
# better safe than sorry
assert len(x.shape) == 3, "x must be (c, x, y)"
if verbose: print("step_size:", step_size)
if verbose: print("do mirror:", do_mirroring)
assert patch_size is not None, "patch_size cannot be None for tiled prediction"
# for sliding window inference the image must at least be as large as the patch size. It does not matter
# whether the shape is divisible by 2**num_pool as long as the patch size is
data, slicer = pad_nd_image(x, patch_size, pad_border_mode, pad_kwargs,
True, None)
data_shape = data.shape # still c, x, y
# compute the steps for sliding window
steps = self._compute_steps_for_sliding_window(patch_size,
data_shape[1:],
step_size)
num_tiles = len(steps[0]) * len(steps[1])
if verbose:
print("data shape:", data_shape)
print("patch size:", patch_size)
print("steps (x, y, and z):", steps)
print("number of tiles:", num_tiles)
# we only need to compute that once. It can take a while to compute this due to the large sigma in
# gaussian_filter
if use_gaussian and num_tiles > 1:
if self._gaussian_2d is None or not all([
i == j for i, j in zip(patch_size,
self._patch_size_for_gaussian_2d)
]):
if verbose: print('computing Gaussian')
gaussian_importance_map = self._get_gaussian(patch_size,
sigma_scale=1. /
8)
self._gaussian_2d = gaussian_importance_map
self._patch_size_for_gaussian_2d = patch_size
else:
if verbose: print("using precomputed Gaussian")
gaussian_importance_map = self._gaussian_2d
gaussian_importance_map = torch.from_numpy(gaussian_importance_map)
if torch.cuda.is_available():
gaussian_importance_map = gaussian_importance_map.cuda(
self.get_device(), non_blocking=True)
else:
gaussian_importance_map = None
if all_in_gpu:
# If we run the inference in GPU only (meaning all tensors are allocated on the GPU, this reduces
# CPU-GPU communication but required more GPU memory) we need to preallocate a few things on GPU
if use_gaussian and num_tiles > 1:
# half precision for the outputs should be good enough. If the outputs here are half, the
# gaussian_importance_map should be as well
gaussian_importance_map = gaussian_importance_map.half()
# make sure we did not round anything to 0
gaussian_importance_map[
gaussian_importance_map == 0] = gaussian_importance_map[
gaussian_importance_map != 0].min()
add_for_nb_of_preds = gaussian_importance_map
else:
add_for_nb_of_preds = torch.ones(patch_size,
device=self.get_device())
if verbose: print("initializing result array (on GPU)")
aggregated_results = torch.zeros([self.num_classes] +
list(data.shape[1:]),
dtype=torch.half,
device=self.get_device())
if verbose: print("moving data to GPU")
data = torch.from_numpy(data).cuda(self.get_device(),
non_blocking=True)
if verbose: print("initializing result_numsamples (on GPU)")
aggregated_nb_of_predictions = torch.zeros(
[self.num_classes] + list(data.shape[1:]),
dtype=torch.half,
device=self.get_device())
else:
if use_gaussian and num_tiles > 1:
add_for_nb_of_preds = self._gaussian_2d
else:
add_for_nb_of_preds = np.ones(patch_size, dtype=np.float32)
aggregated_results = np.zeros([self.num_classes] +
list(data.shape[1:]),
dtype=np.float32)
aggregated_nb_of_predictions = np.zeros([self.num_classes] +
list(data.shape[1:]),
dtype=np.float32)
for x in steps[0]:
lb_x = x
ub_x = x + patch_size[0]
for y in steps[1]:
lb_y = y
ub_y = y + patch_size[1]
predicted_patch = self._internal_maybe_mirror_and_pred_2D(
data[None, :, lb_x:ub_x, lb_y:ub_y], mirror_axes,
do_mirroring, gaussian_importance_map)[0]
if all_in_gpu:
predicted_patch = predicted_patch.half()
else:
predicted_patch = predicted_patch.cpu().numpy()
aggregated_results[:, lb_x:ub_x, lb_y:ub_y] += predicted_patch
aggregated_nb_of_predictions[:, lb_x:ub_x,
lb_y:ub_y] += add_for_nb_of_preds
# we reverse the padding here (remeber that we padded the input to be at least as large as the patch size
slicer = tuple([
slice(0, aggregated_results.shape[i])
for i in range(len(aggregated_results.shape) - (len(slicer) - 1))
] + slicer[1:])
aggregated_results = aggregated_results[slicer]
aggregated_nb_of_predictions = aggregated_nb_of_predictions[slicer]
# computing the class_probabilities by dividing the aggregated result with result_numsamples
class_probabilities = aggregated_results / aggregated_nb_of_predictions
if regions_class_order is None:
predicted_segmentation = class_probabilities.argmax(0)
else:
if all_in_gpu:
class_probabilities_here = class_probabilities.detach().cpu(
).numpy()
else:
class_probabilities_here = class_probabilities
predicted_segmentation = np.zeros(
class_probabilities_here.shape[1:], dtype=np.float32)
for i, c in enumerate(regions_class_order):
predicted_segmentation[class_probabilities_here[i] > 0.5] = c
if all_in_gpu:
if verbose: print("copying results to CPU")
if regions_class_order is None:
predicted_segmentation = predicted_segmentation.detach().cpu(
).numpy()
class_probabilities = class_probabilities.detach().cpu().numpy()
if verbose: print("prediction done")
return predicted_segmentation, class_probabilities
def _internal_predict_3D_tiled(net,
x,
num_repeats=1,
BATCH_SIZE=None,
tile_in_z=True,
step=2,
do_mirroring=True,
mirror_axes=(0, 1, 2),
patch_size=None,
regions_class_order=None,
use_gaussian=False,
pad_border_mode="constant",
pad_kwargs=None,
all_in_gpu=False,
num_classes=None,
return_softmax=True):
assert len(x.shape) == 4, "x must be (c, x, y, z)"
print("step:", step)
print("do mirror:", do_mirroring)
torch.cuda.empty_cache()
with torch.no_grad():
assert patch_size is not None, "patch_size cannot be None for tiled prediction"
data, slicer = pad_nd_image(x, patch_size, pad_border_mode, pad_kwargs,
True, None) # x (1,303,265,372)
data = data[None] # data shape (1,1,303,265,372)
if BATCH_SIZE is not None:
data = np.vstack([data] * BATCH_SIZE)
# input_size = [1, x.shape[0]] + list(patch_size)
if not tile_in_z:
# input_size[2] = data.shape[2]
patch_size[0] = data.shape[2]
# input_size = [int(i) for i in input_size]
if num_classes is None:
num_classes = net.num_classes
if use_gaussian:
global _gaussian_3d, _patch_size_for_gaussian_3d
if _gaussian_3d is None or not all([
i == j
for i, j in zip(patch_size, _patch_size_for_gaussian_3d)
]):
tmp = np.zeros(patch_size)
center_coords = [i // 2 for i in patch_size]
sigmas = [i / 8 for i in patch_size]
tmp[tuple(center_coords)] = 1
tmp_smooth = gaussian_filter(tmp,
sigmas,
0,
mode='constant',
cval=0)
tmp_smooth = tmp_smooth / tmp_smooth.max() * 1
add = tmp_smooth # + 1e-7
# tmp_smooth cannot be 0, otherwise we may end up with nans!
tmp_smooth[tmp_smooth == 0] = np.min(
tmp_smooth[tmp_smooth != 0])
# we only need to compute that once. It can take a while to compute this due to the large sigma in
# gaussian_filter
_gaussian_3d = np.copy(add)
_patch_size_for_gaussian_3d = deepcopy(patch_size)
else:
print("using precomputed Gaussian")
add = _gaussian_3d
else:
add = np.ones(patch_size, dtype=np.float32)
add = add.astype(np.float32)
data_shape = data.shape
print("configuring tiles")
center_coord_start = np.array([i // 2 for i in patch_size]).astype(int)
center_coord_end = np.array([
data_shape[i + 2] - patch_size[i] // 2
for i in range(len(patch_size))
]).astype(int)
num_steps = np.ceil([(center_coord_end[i] - center_coord_start[i]) /
(patch_size[i] / step) for i in range(3)])
step_size = np.array([(center_coord_end[i] - center_coord_start[i]) /
(num_steps[i] + 1e-8) for i in range(3)])
step_size[step_size == 0] = 9999999
xsteps = np.round(
np.arange(center_coord_start[0], center_coord_end[0] + 1e-8,
step_size[0])).astype(int)
ysteps = np.round(
np.arange(center_coord_start[1], center_coord_end[1] + 1e-8,
step_size[1])).astype(int)
zsteps = np.round(
np.arange(center_coord_start[2], center_coord_end[2] + 1e-8,
step_size[2])).astype(int)
if all_in_gpu:
print("initializing result array (on GPU)")
# some of these can remain in half. We just need the reuslts for softmax so it won't hurt at all to reduce
# precision. Inference is of course done in float
result = torch.zeros([num_classes] + list(data.shape[2:]),
dtype=torch.half,
device=get_device(net))
print("moving data to GPU")
data = torch.from_numpy(data).cuda(get_device(net),
non_blocking=True)
print("initializing result_numsamples (on GPU)")
result_numsamples = torch.zeros([num_classes] +
list(data.shape[2:]),
dtype=torch.half,
device=get_device(net))
print("moving add to GPU")
# add = torch.from_numpy(add).cuda(get_device(net), non_blocking=True).half()
add = torch.from_numpy(add).cuda(get_device(net),
non_blocking=True)
add_torch = add
else:
result = np.zeros([num_classes] + list(data.shape[2:]),
dtype=np.float32)
result_numsamples = np.zeros([num_classes] + list(data.shape[2:]),
dtype=np.float32)
add_torch = torch.from_numpy(add).cuda(get_device(net),
non_blocking=True)
print("data shape:", data_shape)
print("patch size:", patch_size)
print("steps (x, y, and z):", xsteps, ysteps, zsteps)
print("number of tiles:", len(xsteps) * len(ysteps) * len(zsteps))
# data, result and add_torch and result_numsamples are now on GPU
for x in xsteps:
lb_x = x - patch_size[0] // 2
ub_x = x + patch_size[0] // 2
for y in ysteps:
lb_y = y - patch_size[1] // 2
ub_y = y + patch_size[1] // 2
for z in zsteps:
lb_z = z - patch_size[2] // 2
ub_z = z + patch_size[2] // 2
predicted_patch = \
_internal_maybe_mirror_and_pred_3D(net, data[:, :, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z],
num_repeats, num_classes, mirror_axes, do_mirroring, add_torch)[0]
if all_in_gpu:
predicted_patch = predicted_patch.half()
else:
predicted_patch = predicted_patch.cpu().numpy()
result[:, lb_x:ub_x, lb_y:ub_y,
lb_z:ub_z] += predicted_patch
if all_in_gpu:
result_numsamples[:, lb_x:ub_x, lb_y:ub_y,
lb_z:ub_z] += add.half()
else:
result_numsamples[:, lb_x:ub_x, lb_y:ub_y,
lb_z:ub_z] += add
slicer = tuple([
slice(0, result.shape[i])
for i in range(len(result.shape) - (len(slicer) - 1))
] + slicer[1:])
result = result[slicer]
result_numsamples = result_numsamples[slicer]
softmax_pred = result / result_numsamples # C,D,H,W
# patient_data = patient_data[:, :old_shape[0], :old_shape[1], :old_shape[2]]
if regions_class_order is None:
predicted_segmentation = softmax_pred.argmax(0)
else:
softmax_pred_here = softmax_pred
predicted_segmentation_shp = softmax_pred_here[0].shape
predicted_segmentation = np.zeros(predicted_segmentation_shp,
dtype=np.float32)
for i, c in enumerate(regions_class_order):
predicted_segmentation[softmax_pred_here[i] > 0.5] = c
if all_in_gpu:
print("copying results to CPU")
predicted_segmentation = predicted_segmentation.detach().cpu(
).numpy()
if return_softmax:
softmax_pred = softmax_pred.half().detach().cpu().numpy()
else:
softmax_pred = None
print("prediction on GPU done")
return predicted_segmentation, softmax_pred
def generate_train_batch(self):
indices = self.get_indices()
imgs = []
labels = []
properties = []
for idx in indices:
item = self._data[idx]
if self.cfg.is_test is False:
img, lbl, _property = item[IMG], item[LABEL], item[PROPERTIES]
stacked_volume = np.stack([img, lbl]) # (2, 512, 512)
assert len(img.shape) == len(
self.cfg.patch_size
), "len(patch_size) must be equal to len(img.shape)"
padded_stacked_volume = pad_nd_image(
stacked_volume, self.cfg.patch_size
) # in case the img_size is smaller than patch_size
padded_stacked_volume = np.expand_dims(padded_stacked_volume,
axis=0) # (1, 2, *size)
if self.cfg.three_dim:
cropped_stacked_volume = random_crop_3D_image_batched(
padded_stacked_volume, self.cfg.patch_size)
else:
cropped_stacked_volume = random_crop_2D_image_batched(
padded_stacked_volume, self.cfg.patch_size)
cropped_stacked_volume = np.squeeze(
cropped_stacked_volume) # (2, *patch_size)
img, lbl = cropped_stacked_volume[0], cropped_stacked_volume[1]
imgs.append(img)
labels.append(lbl)
properties.append(_property)
else:
img, _property = item[IMG], item[PROPERTIES]
assert len(img.shape) == len(
self.cfg.patch_size
), "len(patch_size) must be equal to len(img.shape)"
padded_stacked_volume = pad_nd_image(
img, self.cfg.patch_size
) # in case the img_size is smaller than patch_size
if self.cfg.three_dim:
cropped_stacked_volume = random_crop_3D_image(
padded_stacked_volume, self.cfg.patch_size)
else:
cropped_stacked_volume = random_crop_2D_image(
padded_stacked_volume, self.cfg.patch_size)
imgs.append(cropped_stacked_volume)
properties.append(_property)
batch_img = np.expand_dims(np.stack(imgs),
axis=1) # (b, c, *patch_size)
if self.cfg.is_test:
return {
IMG: batch_img,
BATCH_KEYS: indices,
PROPERTIES: properties
}
batch_label = np.expand_dims(np.stack(labels),
axis=1) # (b, c, *patch_size)
return {
IMG: batch_img,
LABEL: batch_label,
BATCH_KEYS: indices,
PROPERTIES: properties
}
def apply_to_case_uncert(self,
subject,
do_mirroring=[False],
rot_angles=[-15, 0, 15],
rot_axes=[0],
bg_zero=True):
print(
f'applying {len(self.models)} model(s) over {len(self.axes)} axes rotating through {len(rot_angles)} angle(s)'
)
with torch.no_grad():
ensemble_logits = []
ensemble_flips = []
slice_masks = []
case_data = load_subject_and_preprocess(subject,
axis='sagittal',
bg_zero=bg_zero)
if do_mirroring == False:
do_mirroring = [False]
if do_mirroring == True:
do_mirroring = [True, False]
for model in self.models:
model.eval()
for axis in self.axes:
for mirror in do_mirroring:
for angle in rot_angles:
for rot_axis in rot_axes:
if angle != 0:
image_data = rotate_stack(
case_data.copy(), angle, rot_axis)
else:
image_data = case_data.copy()
if mirror:
image_data = image_data[:, ::-1].copy()
if axis == 'coronal':
image_data = np.swapaxes(image_data, 1, 2)
if axis == 'axial':
image_data = np.swapaxes(image_data, 1, 3)
if self.network_type == '3D-to-2D':
input, slicer = pad_nd_image(
image_data, (0, self.patch_size[1],
self.patch_size[2]),
return_slicer=True)
if self.network_type == '3D':
input, slicer = pad_nd_image(
image_data,
self.patch_size,
return_slicer=True)
slicer[0] = slice(
0,
len(self.target_label_sets) * 2, None)
# MSHELLER: input likely needs to be to(device) and not cuda()
output = model.predict_3D(
torch.from_numpy(input).float().cuda(),
do_mirroring=do_mirroring,
patch_size=self.patch_size,
use_sliding_window=True,
use_gaussian=self.use_gaussian)
output = output[1][tuple(slicer)]
slice_sums = np.sum(np.any(image_data > 0, 0),
(1, 2))
slice_mask = np.stack([
np.stack([slice_sums > 2500] *
image_data.shape[2], -1)
] * image_data.shape[3], -1)
slice_mask = np.stack(
[slice_mask] *
len(self.target_label_sets)).astype(
np.uint8)
if axis == 'coronal':
output = np.swapaxes(output, 1, 2)
image_data = np.swapaxes(image_data, 1, 2)
slice_mask = np.swapaxes(slice_mask, 1, 2)
if axis == 'axial':
output = np.swapaxes(output, 1, 3)
image_data = np.swapaxes(image_data, 1, 3)
slice_mask = np.swapaxes(slice_mask, 1, 3)
if mirror:
output = output[:, ::-1].copy()
image_data = image_data[:, ::-1].copy()
slice_mask = slice_mask[:, ::-1].copy()
if angle != 0:
output = rotate_stack(
output.copy(), -angle, rot_axis)
slice_mask = (rotate_stack(
slice_mask, -angle, rot_axis) >
0).astype(np.uint8)
output[:len(self.target_label_names)][
np.logical_not(slice_mask)] = np.nan
ensemble_logits.append(
output[:len(self.target_label_names)])
flip = ((-torch.from_numpy(
output[len(self.target_label_names):]).exp(
)).sigmoid()).numpy()
flip_logit = ((-torch.from_numpy(
output[len(self.target_label_names):]).exp(
))).numpy()
flip[np.logical_not(slice_mask)] = np.nan
ensemble_flips.append(flip)
slice_masks.append(slice_mask)
ensemble_counts = np.sum(slice_masks, axis=0)
uncertainty_weighted_logits = -np.sign(
np.array(ensemble_logits)) * np.array(flip_logit)
full_logit = np.sum(np.divide(np.nan_to_num(np.array(ensemble_logits),
0),
ensemble_counts,
out=np.zeros_like(
np.array(ensemble_logits)),
where=ensemble_counts != 0),
axis=0)
ensemble_predictions = np.greater(np.nan_to_num(ensemble_logits, -10),
0)
full_predictions = np.stack(
[np.greater(full_logit, 0)] *
(len(self.axes) * len(do_mirroring) * len(rot_angles) *
len(rot_axes) * len(self.models)), 0)
preds_agree = np.equal(full_predictions,
ensemble_predictions).astype(np.uint8)
with np.errstate(divide='ignore', invalid='ignore'):
full_flips = np.sum(np.nan_to_num(np.array(ensemble_flips), 0) *
(preds_agree) + np.nan_to_num(
(1 - np.array(ensemble_flips)), 0) *
(1 - preds_agree),
axis=0) / ensemble_counts
full_flips = np.nan_to_num(full_flips, 0)
return np.concatenate([full_logit, full_flips])
def generate_train_batch(self):
selected_keys = np.random.choice(self.list_of_keys, self.batch_size,
True, None)
data = []
seg = []
case_properties = []
for j, i in enumerate(selected_keys):
properties = self._data[i]['properties']
case_properties.append(properties)
if self.get_do_oversample(j):
force_fg = True
else:
force_fg = False
if not isfile(self._data[i]['data_file'][:-4] + ".npy"):
# lets hope you know what you're doing
case_all_data = np.load(self._data[i]['data_file'][:-4] +
".npz")['data']
else:
case_all_data = np.load(
self._data[i]['data_file'][:-4] + ".npy", self.memmap_mode)
# this is for when there is just a 2d slice in case_all_data (2d support)
if len(case_all_data.shape) == 3:
case_all_data = case_all_data[:, None]
if self.transpose is not None:
leading_axis = self.transpose[0]
else:
leading_axis = 0
if not force_fg:
random_slice = np.random.choice(
case_all_data.shape[leading_axis + 1])
else:
# select one class, then select a slice that contains that class
classes_in_slice_per_axis = properties.get(
"classes_in_slice_per_axis")
possible_classes = np.unique(properties['classes'])
possible_classes = possible_classes[possible_classes > 0]
if len(possible_classes) > 0 and not (
0 in possible_classes.shape):
selected_class = np.random.choice(possible_classes)
else:
selected_class = 0
if classes_in_slice_per_axis is not None:
valid_slices = classes_in_slice_per_axis[leading_axis][
selected_class]
else:
valid_slices = np.where(
np.sum(case_all_data[-1] == selected_class,
axis=[i for i in range(3)
if i != leading_axis]))[0]
if len(valid_slices) != 0:
random_slice = np.random.choice(valid_slices)
else:
random_slice = np.random.choice(
case_all_data.shape[leading_axis + 1])
if self.pseudo_3d_slices == 1:
if leading_axis == 0:
case_all_data = case_all_data[:, random_slice]
elif leading_axis == 1:
case_all_data = case_all_data[:, :, random_slice]
else:
case_all_data = case_all_data[:, :, :, random_slice]
if self.transpose is not None and self.transpose[
1] > self.transpose[2]:
case_all_data = case_all_data.transpose(0, 2, 1)
else:
assert leading_axis == 0, "pseudo_3d_slices works only without transpose for now!"
mn = random_slice - (self.pseudo_3d_slices - 1) // 2
mx = random_slice + (self.pseudo_3d_slices - 1) // 2 + 1
valid_mn = max(mn, 0)
valid_mx = min(mx, case_all_data.shape[1])
case_all_seg = case_all_data[-1:]
case_all_data = case_all_data[:-1]
case_all_data = case_all_data[:, valid_mn:valid_mx]
case_all_seg = case_all_seg[:, random_slice]
need_to_pad_below = valid_mn - mn
need_to_pad_above = mx - valid_mx
if need_to_pad_below > 0:
shp_for_pad = np.array(case_all_data.shape)
shp_for_pad[1] = need_to_pad_below
case_all_data = np.concatenate(
(np.zeros(shp_for_pad), case_all_data), 1)
if need_to_pad_above > 0:
shp_for_pad = np.array(case_all_data.shape)
shp_for_pad[1] = need_to_pad_above
case_all_data = np.concatenate(
(case_all_data, np.zeros(shp_for_pad)), 1)
case_all_data = case_all_data.reshape(
(-1, case_all_data.shape[-2], case_all_data.shape[-1]))
case_all_data = np.concatenate((case_all_data, case_all_seg),
0)
num_seg = 1
# why we need this is a little complicated. It has to do with downstream random cropping during data
# augmentation. Basically we will rotate the patch and then to a center crop of size self.final_patch_size.
# Depending on the rotation, scaling and elastic deformation parameters, self.patch_size has to be large
# enough to prevent border artifacts from being present in the final patch
new_shp = None
if np.any(self.need_to_pad) > 0:
new_shp = np.array(case_all_data.shape[1:] + self.need_to_pad)
if np.any(new_shp - np.array(self.patch_size) < 0):
new_shp = np.max(
np.vstack(
(new_shp[None], np.array(self.patch_size)[None])),
0)
else:
if np.any(
np.array(case_all_data.shape[1:]) -
np.array(self.patch_size) < 0):
new_shp = np.max(
np.vstack((np.array(case_all_data.shape[1:])[None],
np.array(self.patch_size)[None])), 0)
if new_shp is not None:
case_all_data_donly = pad_nd_image(case_all_data[:-num_seg],
new_shp,
self.pad_mode,
kwargs=self.pad_kwargs_data)
case_all_data_segnonly = pad_nd_image(
case_all_data[-num_seg:],
new_shp,
'constant',
kwargs={'constant_values': -1})
case_all_data = np.vstack(
(case_all_data_donly, case_all_data_segnonly))[None]
else:
case_all_data = case_all_data[None]
if not force_fg:
case_all_data = random_crop_2D_image_batched(
case_all_data, tuple(self.patch_size))
else:
case_all_data = crop_2D_image_force_fg(case_all_data[0],
tuple(self.patch_size),
selected_class)[None]
data.append(case_all_data[:, :-num_seg])
seg.append(case_all_data[:, -num_seg:])
data = np.vstack(data)
seg = np.vstack(seg)
keys = selected_keys
return {
'data': data,
'seg': seg,
'properties': case_properties,
"keys": keys
}
def generate_train_batch(self):
# DataLoader has its own methods for selecting what patients to use next, see its Documentation
idx = self.get_indices()
patients_for_batch = [self._data[i] for i in idx]
print(idx)
patch_size = None
if self.patch_size == None:
shapes = []
for i in patients_for_batch:
patient_data, _ = self.load_patient(i)
shapes.append(patient_data[0].shape)
patch_size = np.max(shapes, 0)
patch_size = (patch_size[0] + 16 - patch_size[0] % 16,
patch_size[1] + 16 - patch_size[1] % 16,
patch_size[2] + 16 - patch_size[2] % 16)
else:
patch_size = self.patch_size
# initialize empty array for data and seg
data = np.zeros((self.batch_size, self.num_modalities, *patch_size),
dtype=np.float32)
seg = np.zeros((self.batch_size, 1, *patch_size), dtype=np.float32)
metadata = []
patient_names = []
# iterate over patients_for_batch and include them in the batch
for i, j in enumerate(patients_for_batch):
patient_data, patient_metadata = self.load_patient(j)
if self.patch_size is not None:
# this will only pad patient_data if its shape is smaller than self.patch_size
patient_data = pad_nd_image(patient_data, patch_size)
# now random crop to self.patch_size
# crop expects the data to be (b, c, x, y, z) but patient_data is (c, x, y, z) so we need to add one
# dummy dimension in order for it to work (@Todo, could be improved)
patient_data, patient_seg = crop(patient_data[:-1][None],
patient_data[-1:][None],
patch_size,
crop_type="random")
data[i] = patient_data[0]
seg[i] = patient_seg[0]
else:
for j in range(len(patient_data)):
if j != len(patient_data) - 1:
data[i][j] = resize_image_by_padding(
patient_data[j], patch_size)
else:
seg[i][0] = resize_image_by_padding(
patient_data[j], patch_size)
metadata.append(patient_metadata)
patient_names.append(j)
return {
'data': data,
'seg': seg,
'metadata': metadata,
'names': patient_names
}
def _internal_predict_2D_2Dconv_tiled(self,
patient_data,
num_repeats,
BATCH_SIZE=None,
step=2,
do_mirroring=True,
mirror_axes=(0, 1),
patch_size=None,
regions_class_order=None,
use_gaussian=False,
pad_border_mode="edge",
pad_kwargs=None):
with torch.no_grad():
tile_size = patch_size
assert tile_size is not None, "patch_size cannot be None for tiled prediction"
# pad images so that their size is a multiple of tile_size
data, slicer = pad_nd_image(patient_data, tile_size,
pad_border_mode, pad_kwargs, True)
data = data[None]
if BATCH_SIZE is not None:
data = np.vstack([data] * BATCH_SIZE)
input_size = [1, patient_data.shape[0]] + list(tile_size)
input_size = [int(i) for i in input_size]
a = torch.rand(input_size).float()
if self.get_device() == "cpu":
a = a.cpu()
else:
a = a.cuda(self.get_device())
# dummy run to see number of classes
nb_of_classes = self(a).size()[1]
result = np.zeros([nb_of_classes] + list(data.shape[2:]),
dtype=np.float32)
result_numsamples = np.zeros([nb_of_classes] +
list(data.shape[2:]),
dtype=np.float32)
if use_gaussian:
tmp = np.zeros(tile_size, dtype=np.float32)
center_coords = [i // 2 for i in tile_size]
sigmas = [i // 8 for i in tile_size]
tmp[tuple(center_coords)] = 1
tmp_smooth = gaussian_filter(tmp,
sigmas,
0,
mode='constant',
cval=0)
tmp_smooth = tmp_smooth / tmp_smooth.max() * 1
add = tmp_smooth
else:
add = np.ones(tile_size)
if self.get_device() == "cpu":
add_torch = torch.from_numpy(add).cpu().float()
else:
add_torch = torch.from_numpy(add).cuda(
self.get_device()).float()
data_shape = data.shape
center_coord_start = np.array([i // 2
for i in patch_size]).astype(int)
center_coord_end = np.array([
data_shape[i + 2] - patch_size[i] // 2
for i in range(len(patch_size))
]).astype(int)
num_steps = np.ceil([
(center_coord_end[i] - center_coord_start[i]) /
(patch_size[i] / step) for i in range(2)
])
step_size = np.array([
(center_coord_end[i] - center_coord_start[i]) /
(num_steps[i] + 1e-8) for i in range(2)
])
step_size[step_size == 0] = 9999999
xsteps = np.round(
np.arange(center_coord_start[0], center_coord_end[0] + 1e-8,
step_size[0])).astype(int)
ysteps = np.round(
np.arange(center_coord_start[1], center_coord_end[1] + 1e-8,
step_size[1])).astype(int)
for x in xsteps:
lb_x = x - patch_size[0] // 2
ub_x = x + patch_size[0] // 2
for y in ysteps:
lb_y = y - patch_size[1] // 2
ub_y = y + patch_size[1] // 2
result[:, lb_x:ub_x, lb_y:ub_y] += \
self._internal_maybe_mirror_and_pred_2D(data[:, :, lb_x:ub_x, lb_y:ub_y],
num_repeats, mirror_axes, do_mirroring, add_torch)[0]
result_numsamples[:, lb_x:ub_x, lb_y:ub_y] += add
slicer = tuple([
slice(0, result.shape[i])
for i in range(len(result.shape) - (len(slicer) - 1))
] + slicer[1:])
result = result[slicer]
result_numsamples = result_numsamples[slicer]
softmax_pred = result / result_numsamples
if regions_class_order is None:
predicted_segmentation = softmax_pred.argmax(0)
else:
predicted_segmentation_shp = softmax_pred[0].shape
predicted_segmentation = np.zeros(predicted_segmentation_shp,
dtype=np.float32)
for i, c in enumerate(regions_class_order):
predicted_segmentation[softmax_pred[i] > 0.5] = c
return predicted_segmentation, None, softmax_pred, None
def _internal_predict_2D_2Dconv_tiled(self,
patient_data,
num_repeats,
BATCH_SIZE=None,
step=2,
do_mirroring=True,
mirror_axes=(0, 1),
patch_size=None,
regions_class_order=None,
use_gaussian=False,
pad_border_mode="edge",
pad_kwargs=None,
all_in_gpu=False,
nb_of_classes=None):
with torch.no_grad():
tile_size = patch_size
assert tile_size is not None, "patch_size cannot be None for tiled prediction"
# pad images so that their size is a multiple of tile_size
data, slicer = pad_nd_image(patient_data, tile_size,
pad_border_mode, pad_kwargs, True)
if len(data.shape) == 3:
data = data[None]
if BATCH_SIZE is not None:
data = np.vstack([data] * BATCH_SIZE)
else:
assert BATCH_SIZE is None, "cannot overwrite batch size if a batch of data is provided"
if nb_of_classes is None:
input_size = [1, patient_data.shape[0]] + list(tile_size)
input_size = [int(i) for i in input_size]
a = torch.zeros(input_size,
dtype=torch.float).cuda(self.get_device(),
non_blocking=True)
# dummy run to see number of classes
nb_of_classes = self(a).size()[1]
if use_gaussian:
tmp = np.zeros(tile_size, dtype=np.float32)
center_coords = [i // 2 for i in tile_size]
sigmas = [i // 8 for i in tile_size]
tmp[tuple(center_coords)] = 1
tmp_smooth = gaussian_filter(tmp,
sigmas,
0,
mode='constant',
cval=0)
tmp_smooth = tmp_smooth / tmp_smooth.max() * 1
add = tmp_smooth
else:
add = np.ones(tile_size, dtype=np.float32)
add = add.astype(np.float32)
data_shape = data.shape
center_coord_start = np.array([i // 2
for i in patch_size]).astype(int)
center_coord_end = np.array([
data_shape[i + 2] - patch_size[i] // 2
for i in range(len(patch_size))
]).astype(int)
num_steps = np.ceil([
(center_coord_end[i] - center_coord_start[i]) /
(patch_size[i] / step) for i in range(2)
])
step_size = np.array([
(center_coord_end[i] - center_coord_start[i]) /
(num_steps[i] + 1e-8) for i in range(2)
])
step_size[step_size == 0] = 9999999
xsteps = np.round(
np.arange(center_coord_start[0], center_coord_end[0] + 1e-8,
step_size[0])).astype(int)
ysteps = np.round(
np.arange(center_coord_start[1], center_coord_end[1] + 1e-8,
step_size[1])).astype(int)
if all_in_gpu:
# some of these can remain in half. We just need the reuslts for softmax so it won't hurt at all to reduce
# precision. Inference is of course done in float
result = torch.zeros([data.shape[0], nb_of_classes] +
list(data.shape[2:]),
dtype=torch.half,
device=self.get_device())
data = torch.from_numpy(data).cuda(self.get_device(),
non_blocking=True)
result_numsamples = torch.zeros([nb_of_classes] +
list(data.shape[2:]),
dtype=torch.half,
device=self.get_device())
add = torch.from_numpy(add).cuda(self.get_device(),
non_blocking=True)
add_torch = add
else:
result = np.zeros([data.shape[0], nb_of_classes] +
list(data.shape[2:]),
dtype=np.float32)
result_numsamples = np.zeros([nb_of_classes] +
list(data.shape[2:]),
dtype=np.float32)
add_torch = torch.from_numpy(add).cuda(self.get_device(),
non_blocking=True)
for x in xsteps:
lb_x = x - patch_size[0] // 2
ub_x = x + patch_size[0] // 2
for y in ysteps:
lb_y = y - patch_size[1] // 2
ub_y = y + patch_size[1] // 2
predicted_patch = \
self._internal_maybe_mirror_and_pred_2D(data[:, :, lb_x:ub_x, lb_y:ub_y],
num_repeats, mirror_axes, do_mirroring, add_torch)
if all_in_gpu:
predicted_patch = predicted_patch.half()
else:
predicted_patch = predicted_patch.cpu().numpy()
result[:, :, lb_x:ub_x, lb_y:ub_y] += predicted_patch
if all_in_gpu:
result_numsamples[:, lb_x:ub_x,
lb_y:ub_y] += add.half()
else:
result_numsamples[:, lb_x:ub_x, lb_y:ub_y] += add
slicer = tuple([
slice(0, result.shape[i])
for i in range(len(result.shape) - (len(slicer) - 1))
] + slicer[1:])
result = result[slicer]
result_numsamples = result_numsamples[slicer[1:]]
softmax_pred = result[:] / result_numsamples
if regions_class_order is None:
if BATCH_SIZE is not None:
softmax_pred = softmax_pred.mean(0)
predicted_segmentation = softmax_pred.argmax(0)
else:
if all_in_gpu:
softmax_pred_here = softmax_pred.detach().cpu().numpy()
else:
softmax_pred_here = softmax_pred
if BATCH_SIZE is None:
predicted_segmentation_shp = softmax_pred_here[0, 0].shape
predicted_segmentation = np.zeros([
softmax_pred_here.shape[0], *predicted_segmentation_shp
],
dtype=np.float32)
for b in range(softmax_pred_here.shape[0]):
for i, c in enumerate(regions_class_order):
predicted_segmentation[
b, softmax_pred_here[i] > 0.5] = c
else:
predicted_segmentation_shp = softmax_pred_here[0].shape
predicted_segmentation = np.zeros(
predicted_segmentation_shp, dtype=np.float32)
for i, c in enumerate(regions_class_order):
predicted_segmentation[softmax_pred_here[i] > 0.5] = c
if all_in_gpu:
predicted_segmentation = predicted_segmentation.detach().cpu(
).numpy()
softmax_pred = softmax_pred.half().detach().cpu().numpy()
return predicted_segmentation, None, softmax_pred, None
def extract_background(self, data_dict, sample):
"""
Samples background samples with regard to segmentation
Parameters
----------
data_dict : dict
dict with data which should be cropped
sample : dict
additional information which should be passed through to all patches
Needs to provide the segmentation information!
Returns
-------
list of dict
patches
"""
r = np.random.rand(1)
# only generate patches with probability
if not (r < self._prob):
return []
patch_list = []
fseg = sample.pop(self.fseg_key)
fseg[fseg > 0] = 1 # unify all classes
# sample multiple patches
for _ in range(self._num):
iter_count = 0 # count iterations
found_patch = False
# iterate until patch found or max_iter
while not found_patch:
# check for iterations
iter_count += 1
if iter_count >= self.max_iter:
break
slice_list = None
# if it was not possible to extract patch
while slice_list is None:
centre = np.asarray([
np.random.randint(0, i)
for i in data_dict['data'].shape[1:]
])
# extract coordinates
slice_list = LoadPatches.get_patch_coordinates(
centre,
self.patch_size,
data_dict['data'].shape[1:],
random_offset=np.divide(self.patch_size, 4),
min_dist=self.min_dist)
# skip patch if it does not contain enough foreground
slices = [slice(0, fseg.shape[0])] + slice_list
patch_fseg = np.copy(fseg[tuple(slices)])
patcharea = reduce(lambda x, y: x * y, patch_fseg.shape[1:])
if (patch_fseg.sum() / patcharea) < self.fseg_area:
continue
# extract patches and pad data to patch size
patch_dict = {}
for key, item in data_dict.items():
slices = [slice(0, item.shape[0])] + slice_list
patch = np.copy(item[tuple(slices)])
if np.any(patch.shape[1:] < tuple(self.patch_size)):
patch = pad_nd_image(patch,
self.patch_size,
mode='constant',
kwargs={'constant_values': 0})
patch_dict[key] = patch
# discard patches which contain lesions
if np.max(patch_dict[self.mask_key]) > 0:
continue
# generate label
if self.mapping_key is not None:
patch_dict['label'] = np.array(0) # patch label
if not self.pop_fseg:
patch_dict[self.fseg_key] = fseg
found_patch = True
patch_list.append({**patch_dict, **sample})
return patch_list
def _internal_predict_3D_3Dconv_tiled(self,
x,
num_repeats,
BATCH_SIZE=None,
tile_in_z=True,
step=2,
do_mirroring=True,
mirror_axes=(0, 1, 2),
patch_size=None,
regions_class_order=None,
use_gaussian=False,
pad_border_mode="edge",
pad_kwargs=None,
all_in_gpu=False):
"""
x must be (c, x, y, z)
:param x:
:param num_repeats:
:param BATCH_SIZE:
:param tile_in_z:
:param step:
:param do_mirroring:
:param mirror_axes:
:param patch_size:
:param regions_class_order:
:param use_gaussian:
:param pad_border_mode:
:param pad_kwargs:
:param all_in_gpu: if True then data and prediction will be held in GPU for inference. Faster, but uses more vram
:return:
"""
assert len(x.shape) == 4, "x must be (c, x, y, z)"
assert self.get_device() != "cpu"
torch.cuda.empty_cache()
with torch.no_grad():
assert patch_size is not None, "patch_size cannot be None for tiled prediction"
data, slicer = pad_nd_image(x, patch_size, pad_border_mode,
pad_kwargs, True, None)
data = data[None]
if BATCH_SIZE is not None:
data = np.vstack([data] * BATCH_SIZE)
input_size = [1, x.shape[0]] + list(patch_size)
if not tile_in_z:
input_size[2] = data.shape[2]
patch_size[0] = data.shape[2]
input_size = [int(i) for i in input_size]
a = torch.zeros(input_size,
dtype=torch.float).cuda(self.get_device(),
non_blocking=True)
# dummy run to see number of classes
nb_of_classes = self(a).size()[1]
if use_gaussian:
tmp = np.zeros(patch_size)
center_coords = [i // 2 for i in patch_size]
sigmas = [i // 8 for i in patch_size]
tmp[tuple(center_coords)] = 1
tmp_smooth = gaussian_filter(tmp,
sigmas,
0,
mode='constant',
cval=0)
tmp_smooth = tmp_smooth / tmp_smooth.max() * 1
add = tmp_smooth + 1e-8
else:
add = np.ones(patch_size, dtype=np.float32)
add = add.astype(np.float32)
data_shape = data.shape
center_coord_start = np.array([i // 2
for i in patch_size]).astype(int)
center_coord_end = np.array([
data_shape[i + 2] - patch_size[i] // 2
for i in range(len(patch_size))
]).astype(int)
num_steps = np.ceil([
(center_coord_end[i] - center_coord_start[i]) /
(patch_size[i] / step) for i in range(3)
])
step_size = np.array([
(center_coord_end[i] - center_coord_start[i]) /
(num_steps[i] + 1e-8) for i in range(3)
])
step_size[step_size == 0] = 9999999
xsteps = np.round(
np.arange(center_coord_start[0], center_coord_end[0] + 1e-8,
step_size[0])).astype(int)
ysteps = np.round(
np.arange(center_coord_start[1], center_coord_end[1] + 1e-8,
step_size[1])).astype(int)
zsteps = np.round(
np.arange(center_coord_start[2], center_coord_end[2] + 1e-8,
step_size[2])).astype(int)
if all_in_gpu:
# some of these can remain in half. We just need the reuslts for softmax so it won't hurt at all to reduce
# precision. Inference is of course done in float
result = torch.zeros([nb_of_classes] + list(data.shape[2:]),
dtype=torch.half).cuda()
data = torch.from_numpy(data).cuda(self.get_device())
result_numsamples = torch.zeros([nb_of_classes] +
list(data.shape[2:]),
dtype=torch.half).cuda()
add = torch.from_numpy(add).cuda(self.get_device()).float()
add_torch = add
else:
result = np.zeros([nb_of_classes] + list(data.shape[2:]),
dtype=np.float32)
result_numsamples = np.zeros([nb_of_classes] +
list(data.shape[2:]),
dtype=np.float32)
add_torch = torch.from_numpy(add).cuda(self.get_device(),
non_blocking=True)
# data, result and add_torch and result_numsamples are now on GPU
for x in xsteps:
lb_x = x - patch_size[0] // 2
ub_x = x + patch_size[0] // 2
for y in ysteps:
lb_y = y - patch_size[1] // 2
ub_y = y + patch_size[1] // 2
for z in zsteps:
lb_z = z - patch_size[2] // 2
ub_z = z + patch_size[2] // 2
predicted_patch = \
self._internal_maybe_mirror_and_pred_3D(data[:, :, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z],
num_repeats, mirror_axes, do_mirroring, add_torch)[0]
if all_in_gpu:
predicted_patch = predicted_patch.half()
else:
predicted_patch = predicted_patch.cpu().numpy()
result[:, lb_x:ub_x, lb_y:ub_y,
lb_z:ub_z] += predicted_patch
if all_in_gpu:
result_numsamples[:, lb_x:ub_x, lb_y:ub_y,
lb_z:ub_z] += add.half()
else:
result_numsamples[:, lb_x:ub_x, lb_y:ub_y,
lb_z:ub_z] += add
slicer = tuple([
slice(0, result.shape[i])
for i in range(len(result.shape) - (len(slicer) - 1))
] + slicer[1:])
result = result[slicer]
result_numsamples = result_numsamples[slicer]
softmax_pred = result / result_numsamples
# patient_data = patient_data[:, :old_shape[0], :old_shape[1], :old_shape[2]]
if regions_class_order is None:
predicted_segmentation = softmax_pred.argmax(0)
else:
if all_in_gpu:
softmax_pred_here = softmax_pred.detach().cpu().numpy()
else:
softmax_pred_here = softmax_pred
predicted_segmentation_shp = softmax_pred_here[0].shape
predicted_segmentation = np.zeros(predicted_segmentation_shp,
dtype=np.float32)
for i, c in enumerate(regions_class_order):
predicted_segmentation[softmax_pred_here[i] > 0.5] = c
if all_in_gpu:
predicted_segmentation = predicted_segmentation.detach().cpu(
).numpy()
softmax_pred = softmax_pred.half().detach().cpu().numpy()
return predicted_segmentation, None, softmax_pred, None
def extract_patch(self, data_dict, sample):
"""
Samples patches around foreground objects
Parameters
----------
data_dict : dict
dict with data which should be cropped
sample : dict
additional information added to every sample
Returns
-------
list of dict
patches
"""
patch_list = []
# only sample if at least one lesion is present
if data_dict[self.mask_key].max() > 0:
centre = [
i['centroid'] for i in regionprops(data_dict[self.mask_key][0])
]
for c in centre:
random_offset = np.divide(self.patch_size, 4) if \
self.rand_offset else None
# extract coordinates
slice_list = LoadPatches.get_patch_coordinates(
c,
self.patch_size,
data_dict[self.mask_key].shape[1:],
random_offset=random_offset,
min_dist=self.min_dist)
# check if extraction was successful
if slice_list is None:
continue
# extract patches and pad data to patch size
patch_dict = {}
for key, item in data_dict.items():
# use all channels
slices = [slice(0, item.shape[0])] + slice_list
patch = np.copy(item[tuple(slices)])
if np.any(patch.shape[1:] < tuple(self.patch_size)):
patch = pad_nd_image(patch,
self.patch_size,
mode='constant',
kwargs={'constant_values': 0})
patch_dict[key] = patch
# skip patches without background
if self.skip_no_background:
if not (patch_dict[self.mask_key] == 0).any():
continue
# skip samples with multiple foreground objects (>2, 1 bg, 1 fg)
if self.discard_mul_fg and \
np.unique(patch_dict[self.mask_key]).shape[0] > 2:
continue
# generate label for patch
if self.mapping_key is not None:
patch_values = np.unique(patch_dict[self.mask_key])
patch_values = patch_values[patch_values > 0]
mapping = sample[self.mapping_key]
patch_classes = [
mapping[mapping[:, 0] == i][0, 1] for i in patch_values
if (mapping[:, 0] == i).any()
]
if len(patch_classes) > 0:
patch_dict['label'] = np.array(np.max(patch_classes))
else:
if 'id' not in sample:
logger.warning(f'Skip: no patch classes found.')
else:
logger.warning(
f'Skip: {sample["id"]} no patch classes found.'
)
print("Mapping: {}".format(mapping))
print("Val: {}".format(patch_values))
continue
patch_list.append({**patch_dict, **sample})
return patch_list