def test_remove_halo(self): patch_halo = (4, 4, 4) shape = (128, 128, 128) input = np.random.randint(0, 10, size=(1, 16, 16, 16)) index = (slice(0, 1), slice(12, 28), slice(16, 32), slice(16, 32)) u_patch, u_index = remove_halo(input, index, shape, patch_halo) assert np.array_equal(input[:, 4:12, 4:12, 4:12], u_patch) assert u_index == (slice(0, 1), slice(16, 24), slice(20, 28), slice(20, 28)) index = (slice(0, 1), slice(112, 128), slice(112, 128), slice(112, 128)) u_patch, u_index = remove_halo(input, index, shape, patch_halo) assert np.array_equal(input[:, 4:16, 4:16, 4:16], u_patch) assert u_index == (slice(0, 1), slice(116, 128), slice(116, 128), slice(116, 128))
def predict(self): out_channels = self.config['model'].get('out_channels') if out_channels is None: out_channels = self.config['model']['dt_out_channels'] prediction_channel = self.config.get('prediction_channel', None) if prediction_channel is not None: logger.info( f"Using only channel '{prediction_channel}' from the network output" ) device = self.config['device'] output_heads = self.config['model'].get('output_heads', 1) logger.info(f'Running prediction on {len(self.loader)} batches...') # dimensionality of the the output predictions volume_shape = self._volume_shape(self.loader.dataset) if prediction_channel is None: prediction_maps_shape = (out_channels, ) + volume_shape else: # single channel prediction map prediction_maps_shape = (1, ) + volume_shape logger.info( f'The shape of the output prediction maps (CDHW): {prediction_maps_shape}' ) patch_halo = self.predictor_config.get('patch_halo', (4, 8, 8)) self._validate_halo(patch_halo, self.config['loaders']['test']['slice_builder']) logger.info(f'Using patch_halo: {patch_halo}') # create destination H5 file h5_output_file = h5py.File(self.output_file, 'w') # allocate prediction and normalization arrays logger.info('Allocating prediction and normalization arrays...') prediction_maps, normalization_masks = self._allocate_prediction_maps( prediction_maps_shape, output_heads, h5_output_file) # Sets the module in evaluation mode explicitly (necessary for batchnorm/dropout layers if present) self.model.eval() # Set the `testing=true` flag otherwise the final Softmax/Sigmoid won't be applied! self.model.testing = True # Run predictions on the entire input dataset with torch.no_grad(): for batch, indices in self.loader: # send batch to device batch = batch.to(device) # forward pass predictions = self.model(batch) # wrap predictions into a list if there is only one output head from the network if output_heads == 1: predictions = [predictions] # for each output head for prediction, prediction_map, normalization_mask in zip( predictions, prediction_maps, normalization_masks): # convert to numpy array prediction = prediction.cpu().numpy() # for each batch sample for pred, index in zip(prediction, indices): # save patch index: (C,D,H,W) if prediction_channel is None: channel_slice = slice(0, out_channels) else: channel_slice = slice(0, 1) index = (channel_slice, ) + index if prediction_channel is not None: # use only the 'prediction_channel' logger.info( f"Using channel '{prediction_channel}'...") pred = np.expand_dims(pred[prediction_channel], axis=0) logger.info(f'Saving predictions for slice:{index}...') # remove halo in order to avoid block artifacts in the output probability maps u_prediction, u_index = remove_halo( pred, index, volume_shape, patch_halo) # accumulate probabilities into the output prediction array prediction_map[u_index] += u_prediction # count voxel visits for normalization normalization_mask[u_index] += 1 # save results to self._save_results(prediction_maps, normalization_masks, output_heads, h5_output_file, self.loader.dataset) # close the output H5 file h5_output_file.close()