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()