def test_extract_largest_foreground_connected_component( restrict_classes: Optional[List[int]], has_fg_component: bool) -> None: segmentation = np.zeros((4, 4, 4)) if has_fg_component: segmentation[0, :, :] = 1 segmentation[2, 2, 2] = 1 segmentation[1, :, :] = 2 segmentation[3, 3, 3] = 2 # extract the largest components for all foreground classes restrict_pairs = None if restrict_classes is not None: restrict_pairs = [(idx, None) for idx in restrict_classes] largest_fg_classes = image_util.extract_largest_foreground_connected_component( multi_label_array=segmentation, restrictions=restrict_pairs) # type: ignore expected = segmentation.copy() if restrict_classes is None or restrict_classes == [1, 2]: expected[2, 2, 2] = 0 expected[3, 3, 3] = 0 elif restrict_classes == [1]: expected[2, 2, 2] = 0 elif restrict_classes == [2]: expected[3, 3, 3] = 0 assert np.array_equal(largest_fg_classes, expected)
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_extract_largest_foreground_connected_component_with_threshold( ) -> None: # Class 1 is 17 voxels in two components, of 16 voxels and 1 voxel (so latter is < 0.1 as a proportion) # Class 2 is 20 voxels in two components, of 16 voxels and 4 voxel (so latter is > 0.1 as a proportion) segmentation = np.zeros((4, 4, 4)) segmentation[0, :, :] = 1 # 16 voxels segmentation[2, 2, 2] = 1 # 1 voxel segmentation[1, :, :] = 2 # 16 voxels segmentation[3, 3, :] = 2 # 4 voxels # extract the largest components for all foreground classes. A threshold # of 0.1 is greater than 1/17 but not greater than 4/20, so only the lone voxel # at 2,2,2 should be deleted. largest_fg_classes = image_util.extract_largest_foreground_connected_component( multi_label_array=segmentation, restrictions=[(1, 0.1), (2, 0.1)]) expected = segmentation.copy() # We expected only the 1 disconnected voxel in class 1 to disappear, not the 4 in class 2. expected[2, 2, 2] = 0 assert np.array_equal(largest_fg_classes, expected)
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)