def test_aggregate_results() -> None:
"""
Test to make sure inference results are aggregated as expected
"""
torch.manual_seed(1)
num_models = 3
# set expected posteriors
model_results = []
# create results for each model
for x in range(num_models):
posteriors = torch.nn.functional.softmax(torch.rand(3, 3, 3, 3),
dim=0).numpy()
model_results.append(
InferencePipeline.Result(
epoch=0,
patient_id=0,
posteriors=posteriors,
segmentation=posteriors_to_segmentation(posteriors),
voxel_spacing_mm=(1, 1, 1)))
# We calculate expected_posteriors before aggregating, as aggregation modifies model_results.
expected_posteriors = np.mean([x.posteriors for x in model_results],
axis=0)
ensemble_result = EnsemblePipeline.aggregate_results(
model_results, aggregation_type=EnsembleAggregationType.Average)
assert ensemble_result.epoch == model_results[0].epoch
assert ensemble_result.patient_id == model_results[0].patient_id
assert np.array_equal(ensemble_result.posteriors, expected_posteriors)
assert np.array_equal(ensemble_result.segmentation,
posteriors_to_segmentation(expected_posteriors))
def test_posteriors_to_segmentation() -> None:
"""
Test to make sure the posterior to segmentation conversion is as expected
"""
with pytest.raises(ValueError):
# noinspection PyTypeChecker
image_util.posteriors_to_segmentation(posteriors=None)
with pytest.raises(ValueError):
image_util.posteriors_to_segmentation(posteriors=np.zeros(shape=(3, 3,
3)))
with pytest.raises(ValueError):
image_util.posteriors_to_segmentation(posteriors=np.zeros(shape=(3, 3,
3, 3,
3,
3)))
image = np.zeros(shape=(2, 3, 3, 3))
image[1] = 0.6
# class 1 dominates entire image
assert np.array_equal(np.ones(shape=(3, 3, 3)),
image_util.posteriors_to_segmentation(image))
# no dominating class (first index with this argmax chosen chosen by default)
image[...] = 0.5
assert np.array_equal(np.zeros(shape=(3, 3, 3)),
image_util.posteriors_to_segmentation(image))
def post_process(self) -> InferenceBatch:
"""
Perform post processing on the computed outputs of the a single pass of the pipelines.
Currently the following operations are performed:
-------------------------------------------------------------------------------------
1) the mask is applied to the posteriors (if required).
2) the final posteriors are used to perform an argmax to generate a multi-label segmentation.
3) extract the largest foreground connected component in the segmentation if required
"""
mask = self.get_component(InferenceBatch.Components.Mask)
posteriors = self.get_component(InferenceBatch.Components.Posteriors)
if mask is not None:
posteriors = image_util.apply_mask_to_posteriors(
posteriors=posteriors, mask=mask)
# create segmentation using an argmax over the posterior probabilities
segmentation = image_util.posteriors_to_segmentation(posteriors)
# update the post-processed posteriors and save the segmentation
self.set_component(component=InferenceBatch.Components.Posteriors,
data=posteriors)
self.set_component(component=InferenceBatch.Components.Segmentation,
data=segmentation)
return self
def aggregate_results(results: Iterable[InferencePipeline.Result],
aggregation_type: EnsembleAggregationType) -> InferencePipeline.Result:
"""
Helper method to aggregate results from multiple inference pipelines, based on the aggregation type provided.
:param results: inference pipeline results to aggregate. This may be a Generator to prevent multiple large
posterior arrays being held at the same time. The first element of the sequence is modified in place to
minimize memory use.
:param aggregation_type: aggregation function to use to combine the results.
:return: InferenceResult: contains a Segmentation for each of the classes and their posterior
probabilities.
"""
if aggregation_type != EnsembleAggregationType.Average:
raise NotImplementedError(f"Ensembling is not implemented for aggregation type: {aggregation_type}")
aggregate: Optional[InferencePipeline.Result] = None
n_results = 0
for result in results:
if aggregate is None:
aggregate = result
else:
aggregate.posteriors += result.posteriors
n_results += 1
assert aggregate is not None
aggregate.posteriors /= n_results
aggregate.segmentation = posteriors_to_segmentation(aggregate.posteriors)
return aggregate
def test_apply_summed_probability_rules_no_change(
model_config: SegmentationModelBase, is_batched: bool) -> None:
"""
Test `apply_summed_probability_rules` with valid inputs and no expected change
"""
posteriors = np.zeros(shape=(2, 4, 3, 3, 3))
posteriors[:, 3] = 0.6
posteriors[:, 1, :2] = 0.2
posteriors[:, 2, :2] = 0.2
# class 1 and class 2 together do not have a larger probability than external (class 3).
expected_segmentation = np.full(shape=(2, 3, 3, 3), fill_value=3)
if not is_batched:
posteriors = posteriors[0]
expected_segmentation = expected_segmentation[0]
segmentation = image_util.posteriors_to_segmentation(posteriors)
# test for both np arrays and tensors
# we make a copy of segmentation as apply_summed_probability_rules modifies it in place
assert np.array_equal(
expected_segmentation,
image_util.apply_summed_probability_rules(model_config, posteriors,
np.copy(segmentation)))
# noinspection PyTypeChecker
assert torch.equal(
torch.from_numpy(expected_segmentation).type(
torch.LongTensor), # type: ignore
image_util.apply_summed_probability_rules(
model_config, torch.from_numpy(posteriors),
torch.from_numpy(np.copy(segmentation))))
def test_apply_summed_probability_rules_incorrect_input(
model_config: SegmentationModelBase) -> None:
"""
Test `apply_summed_probability_rules` with invalid inputs
"""
posteriors = np.zeros(shape=(2, 4, 3, 3, 3))
segmentation = image_util.posteriors_to_segmentation(posteriors)
with pytest.raises(ValueError):
# noinspection PyTypeChecker
image_util.apply_summed_probability_rules(model_config,
posteriors=posteriors,
segmentation=None)
with pytest.raises(ValueError):
# noinspection PyTypeChecker
image_util.apply_summed_probability_rules(model_config,
posteriors=None,
segmentation=segmentation)
with pytest.raises(ValueError):
image_util.apply_summed_probability_rules(
model_config,
posteriors=posteriors,
segmentation=np.zeros(shape=(3, 3, 3)))
with pytest.raises(ValueError):
image_util.apply_summed_probability_rules(
model_config,
posteriors=posteriors,
segmentation=np.zeros(shape=(3, 3, 3, 3)))
def post_process(
self,
results: InferencePipeline.Result) -> InferencePipeline.Result:
"""
Perform connected component analysis to update segmentation with largest
connected component based on the configurations
:param results: inference results to post-process
:return: post-processed version of results
"""
if self.model_config.posterior_smoothing_mm:
posteriors = gaussian_smooth_posteriors(
posteriors=results.posteriors,
kernel_size_mm=self.model_config.posterior_smoothing_mm,
voxel_spacing_mm=results.voxel_spacing_mm)
results = InferencePipeline.Result(
epoch=results.epoch,
patient_id=results.patient_id,
posteriors=posteriors,
segmentation=posteriors_to_segmentation(posteriors),
voxel_spacing_mm=results.voxel_spacing_mm)
if self.model_config.summed_probability_rules and not self.model_config.disable_extra_postprocessing:
assert isinstance(self.model_config, SegmentationModelBase)
results = results.with_new_segmentation(
image_util.apply_summed_probability_rules(
self.model_config, results.posteriors,
results.segmentation))
if self.model_config.largest_connected_component_foreground_classes is not None:
# get indices for classes to restrict
restrict_class_indices_and_thresholds = []
for name, idx in self.model_config.class_and_index_with_background(
).items():
for name2, threshold in self.model_config.largest_connected_component_foreground_classes:
if name2 == name:
restrict_class_indices_and_thresholds.append(
(idx, threshold))
results = results.with_new_segmentation(
image_util.extract_largest_foreground_connected_component(
multi_label_array=results.segmentation,
# mypy gets confused below because List is invariant. Sequence is covariant
# but does not allow "append".
restrictions=restrict_class_indices_and_thresholds)
) # type: ignore
if self.model_config.slice_exclusion_rules and not self.model_config.disable_extra_postprocessing:
results = results.with_new_segmentation(
image_util.apply_slice_exclusion_rules(self.model_config,
results.segmentation))
return results
def _forward_pass(
self,
patches: torch.Tensor,
mask: torch.Tensor = None,
labels: torch.Tensor = None) -> SegmentationForwardPass.Result:
# ensure that we always have float tensors as the model is defined over floats
# and transfer the tensors to the GPU if possible before the forward pass
patches = self.config.get_gpu_tensor_if_possible(patches)
if mask is not None:
mask = self.config.get_gpu_tensor_if_possible(mask)
logits, loss = self._compute_loss(patches, labels)
if self.in_training_mode:
if loss is None:
raise ValueError(
"When running training, the labels must be present for loss computation."
)
assert self.optimizer is not None # for mypy
single_optimizer_step(loss, self.optimizer, self.gradient_scaler)
# Aggregate data parallel logits if multiple hardware are used in forward pass
if isinstance(logits, list):
# When using multiple GPUs, logits is a list of tensors. Gather will concatenate them
# across the first dimension, and move them to GPU0.
logits = torch.nn.parallel.gather(logits, target_device=0)
# apply Softmax on dimension 1 (Class) to map model output into a posterior probability distribution [0,1]
posteriors = torch.nn.functional.softmax(logits, dim=1)
# apply mask if required
if not self.in_training_mode and mask is not None:
posteriors = image_util.apply_mask_to_posteriors(
posteriors=posteriors, mask=mask)
# post process posteriors to compute result
segmentations = image_util.posteriors_to_segmentation(
posteriors=posteriors).data.cpu().numpy()
posteriors = posteriors.data.cpu().numpy()
return SegmentationForwardPass.Result(
posteriors=posteriors,
segmentations=segmentations,
loss=loss.item() if loss is not None else None)
def post_process_posteriors(
self,
posteriors: np.ndarray,
mask: np.ndarray = None) -> Tuple[np.ndarray, np.ndarray]:
"""
Perform post processing on the computed outputs of the a single pass of the pipelines.
Currently the following operations are performed:
-------------------------------------------------------------------------------------
1) the mask is applied to the posteriors (if required).
2) the final posteriors are used to perform an argmax to generate a multi-label segmentation.
3) extract the largest foreground connected component in the segmentation if required
"""
if mask is not None:
posteriors = image_util.apply_mask_to_posteriors(
posteriors=posteriors, mask=mask)
# create segmentation using an argmax over the posterior probabilities
segmentation = image_util.posteriors_to_segmentation(posteriors)
return posteriors, segmentation
def training_or_validation_step(self, sample: Dict[str,
Any], batch_index: int,
is_training: bool) -> torch.Tensor:
"""
Runs training for a single minibatch of training or validation data, and computes all metrics.
:param is_training: If true, the method is called from `training_step`, otherwise it is called from
`validation_step`.
:param sample: The batched sample on which the model should be trained.
:param batch_index: The index of the present batch (supplied only for diagnostics).
"""
cropped_sample: CroppedSample = CroppedSample.from_dict(sample=sample)
# Forward propagation can lead to a model output that is smaller than the input image (crop).
# labels_center_crop is the relevant part of the labels tensor that the model will actually produce.
labels = cropped_sample.labels_center_crop
mask = cropped_sample.mask_center_crop if is_training else None
if is_training:
logits = self.model(cropped_sample.image)
else:
with torch.no_grad():
logits = self.model(cropped_sample.image)
loss = self.loss_fn(logits, labels)
# apply Softmax on dimension 1 (Class) to map model output into a posterior probability distribution [0,1]
posteriors = self.logits_to_posterior(logits)
# apply mask if required
if mask is not None:
posteriors = image_util.apply_mask_to_posteriors(
posteriors=posteriors, mask=mask) # type: ignore
# post process posteriors to compute result
segmentation = image_util.posteriors_to_segmentation(
posteriors=posteriors) # type: ignore
self.compute_metrics(cropped_sample, segmentation,
is_training) # type: ignore
self.write_loss(is_training, loss)
return loss
def test_apply_summed_probability_rules_change(
model_config: SegmentationModelBase, is_batched: bool) -> None:
"""
Test `apply_summed_probability_rules` with valid inputs and an expected change
"""
posteriors = np.zeros(shape=(2, 4, 3, 3, 3))
posteriors[:, 3] = 0.4
posteriors[:, 1, :1] = 0.35
posteriors[:, 2, :1] = 0.25
posteriors[:, 1, 1:2] = 0.25
posteriors[:, 2, 1:2] = 0.35
# class 1 and class 2 together have a larger probability than class 3 in half the image
# In this region, we replace external pixels (class 3) with class 1 or class 2, whichever has the higher probability
expected_segmentation = np.full(shape=(2, 3, 3, 3), fill_value=3)
expected_segmentation[:, :1] = 1
expected_segmentation[:, 1:2] = 2
if not is_batched:
posteriors = posteriors[0]
expected_segmentation = expected_segmentation[0]
segmentation = image_util.posteriors_to_segmentation(posteriors)
# test for both np arrays and tensors
# we make a copy of segmentation as apply_summed_probability_rules modifies it in place
assert np.array_equal(
expected_segmentation,
image_util.apply_summed_probability_rules(model_config, posteriors,
np.copy(segmentation)))
# noinspection PyTypeChecker
assert torch.equal(
torch.from_numpy(expected_segmentation).type(
torch.LongTensor), # type: ignore
image_util.apply_summed_probability_rules(
model_config, torch.from_numpy(posteriors),
torch.from_numpy(np.copy(segmentation))))
def inference_identity(
test_output_dirs: OutputFolderForTests,
image_size: Any = (4, 5, 8),
crop_size: Any = (5, 5, 5),
shrink_by: Any = (0, 0, 0),
num_classes: int = 5,
create_mask: bool = True,
extract_largest_foreground_connected_component: bool = False,
is_ensemble: bool = False,
posterior_smoothing_mm: Any = None) -> None:
"""
Test to make sure inference pipeline is identity preserving, ie: we can recreate deterministic
model output, ensuring the patching and stitching is robust.
"""
# fix random seed
np.random.seed(0)
ground_truth_ids = list(map(str, range(num_classes)))
# image to run inference on: The mock model passes the input image through, hence the input
# image must have as many channels as we have classes (plus background), such that the output is
# also a valid posterior.
num_channels = num_classes + 1
image_channels = np.random.randn(num_channels, *list(image_size))
# create a random mask if required
mask = np.round(np.random.uniform(
size=image_size)).astype(np.int) if create_mask else None
config = InferenceIdentityModel(shrink_by=shrink_by)
config.crop_size = crop_size
config.test_crop_size = crop_size
config.image_channels = list(map(str, range(num_channels)))
config.ground_truth_ids = ground_truth_ids
config.posterior_smoothing_mm = posterior_smoothing_mm
# We have to set largest_connected_component_foreground_classes after creating the model config,
# because this parameter is not overridable and hence will not be set by GenericConfig's constructor.
if extract_largest_foreground_connected_component:
config.largest_connected_component_foreground_classes = [
(c, None) for c in ground_truth_ids
]
# set expected posteriors
expected_posteriors = torch.nn.functional.softmax(
torch.tensor(image_channels), dim=0).numpy()
# apply the mask if required
if mask is not None:
expected_posteriors = image_util.apply_mask_to_posteriors(
expected_posteriors, mask)
if posterior_smoothing_mm is not None:
expected_posteriors = image_util.gaussian_smooth_posteriors(
posteriors=expected_posteriors,
kernel_size_mm=posterior_smoothing_mm,
voxel_spacing_mm=(1, 1, 1))
# compute expected segmentation
expected_segmentation = image_util.posteriors_to_segmentation(
expected_posteriors)
if extract_largest_foreground_connected_component:
largest_component = image_util.extract_largest_foreground_connected_component(
multi_label_array=expected_segmentation)
# make sure the test data is accurate by checking if more than one component exists
assert not np.array_equal(largest_component, expected_segmentation)
expected_segmentation = largest_component
# instantiate the model
checkpoint = test_output_dirs.root_dir / "checkpoint.ckpt"
create_model_and_store_checkpoint(config, checkpoint_path=checkpoint)
# create single or ensemble inference pipeline
inference_pipeline = InferencePipeline.create_from_checkpoint(
path_to_checkpoint=checkpoint, model_config=config)
assert inference_pipeline is not None
full_image_inference_pipeline = EnsemblePipeline([inference_pipeline], config) \
if is_ensemble else inference_pipeline
# compute full image inference results
inference_result = full_image_inference_pipeline \
.predict_and_post_process_whole_image(image_channels=image_channels, mask=mask, voxel_spacing_mm=(1, 1, 1))
# Segmentation must have same size as input image
assert inference_result.segmentation.shape == image_size
assert inference_result.posteriors.shape == (num_classes +
1, ) + image_size
# check that the posteriors and segmentations are as expected. Flatten to a list so that the error
# messages are more informative.
assert np.allclose(inference_result.posteriors, expected_posteriors)
assert np.array_equal(inference_result.segmentation, expected_segmentation)
