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