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