Subcortical Segmentation

This module contains the main functions for performing subcortical segmentations on unscaled aligned scans. As post-processing step, only the largest connected component of each class is kept by default. Doing this post-processing step can be disabled by setting the connected_component argument to False.

A single scan can be segmented using the segment_scan_subc() function, which is a wrapper that loads the segmentation model, prepares the scan for segmentation and predicts the segmentation mask.

>>> aligned_scan = torch.rand((160, 160, 160))
>>> segm_pred_np, key_maps = segment_scan_subc(aligned_scan, connected_component=True)

As for the alignment, it is recommended to use the functions load_segmentation_model(), prepare_scan_segm(), and segment_subcortical() directly to have more control over the workflow. This can for example be used as follows:

>>> segm_model = load_segmentation_model()
>>> aligned_scan = torch.rand((160, 160, 160))
>>> aligned_scan_prep = prepare_scan(aligned_scan)
>>> segm_pred, key_maps = segment_subcortical(aligned_scan_prep, segm_model)
>>> segm_pred_cc = keep_largest_compoment(segm_pred)

The segment_subcortical() function can also process batches of data (i.e. multiple scans at once), which can be useful to speed up analysis. More advanced examples can be found in the Example Gallery.

Lastly, the compute_volume_segm() function can be used to compute the volume of each structure in the segmentation mask:

>>> key_maps = {"ChP": 1, "LPVH": 2, "CSPV": 3, "CB": 4}
>>> volume_dict = compute_volume_segm(segm_pred_cc, key_maps, spacing=(0.6, 0.6, 0.6))

Module functions

fetalbrain.structural_segmentation.subcortical_segm.compute_volume_segm(segm_map: ndarray, key_maps: dict[str, int], spacing: tuple[float, float, float]) → dict[str, list]

Get the volume of each structure in a segmentation map.

Parameters:

• segm_map – the segmentation map of size [B, H, W, D] or [H,W,D] with values corresponding to the keys/values in key_maps.

• key_maps – dictionary with the classes as keys, and the integer value in segm_map as value: {class: value}

• spacing – the spacing of the image in mm, should be a tuple of size 3: (x, y, z)

Returns:

volume_dict – dictionary with all classes as keys, and a list of the volume of each structure in cm^3 as value.

Example

>>> segm_map = np.random.randint(0, 5, (1, 160, 160, 160))
>>> key_maps = {"ChP": 1, "LPVH": 2, "CSPV": 3, "CB": 4}
>>> spacing = (0.6, 0.6, 0.6)
>>> volume_dict = compute_volume_segm(segm_map, key_maps, spacing)

fetalbrain.structural_segmentation.subcortical_segm.keep_largest_compoment(segm_map: ndarray) → ndarray

Only keeps the largest connected component for each class

Parameters:

segm_map – input segmentation map of size [B, H, W, D] or [H, W, D] with integers values corresponding to a class

Returns:

conn_comp_segm – segmentation map of size [B, H, W, D] or [H, W, D] with for each class only the largest connected component is kept. The batch dimension is included only if it is larger than 1.

Example

>>> segm_map = np.random.randint(0, 5, (1, 160, 160, 160))
>>> conn_comp_segm = keep_largest_compoment(segm_map)

fetalbrain.structural_segmentation.subcortical_segm.load_segmentation_model(model_path: Path | None = None, device: str = 'cpu') → DataParallel[UNet]

Load the trained segmentation model :param model_path: path to the trained model. Defaults to None (uses the default model).

Returns:

model – segmentation model with trained weights loaded

Example

>>> model = load_segmentation_model()

fetalbrain.structural_segmentation.subcortical_segm.segment_scan_subc(aligned_scan: Tensor | ndarray, connected_component: bool = True) → tuple[ndarray, dict]

full function to segment a single scan or a batch of scans

Parameters:

• aligned_scan – input array or tensor to segment, can be of size [B,H,W,D], [B,C,H,W,D] or [H,W,D]

• connected_component – whether to only keep the largest connected component. Defaults to True.

Returns:

• segm_pred_np – segmentation map of size [B, H, W, D] or [H, W, D] with a multi-class segmentation. The batch dimension is only kepf it is larger than 1.

• key_maps – dictionary with the classes as keys, and the integer value in segm_map as value: {class: value}

Example

>>> aligned_scan = torch.rand((160, 160, 160))
>>> segm_pred_np, key_maps = segment_scan_subc(aligned_scan)

fetalbrain.structural_segmentation.subcortical_segm.segment_subcortical(aligned_scan: Tensor, segm_model: DataParallel[UNet], connected_component: bool = False) → tuple[ndarray, dict[str, int]]

Generate subcortical predictions for a given aligned scan using the provided segmentation model.

Parameters:

• aligned_scan – torch tensor containing the aligned image(s) to segment. Should be of size [B, 1, H, W, D]. Expected input orientation is equal to the output of the align_to_atlas(scale = False) function.

• segm_model – a loaded segmentation model, can be obtained using the load_segmentation_model() function.

Returns:

• segm_map – multiclass segmentation map of size [B, H, W, D] with values between 0 and 4. The values correspond to the following classes: 0 = background, 1 = choroid plexus (ChP), 2 = lateral posterior ventricle horn (LPVH), 3 = cavum septum pellucidum et vergae (CSPV), 4 = cerebellum (CB).

• key_maps – dictionary containing the mapping between the class values and the class names.

Example

>>> aligned_scan = torch.rand((160, 160, 160))
>>> aligned_scan_prep = prepare_scan(aligned_scan)
>>> segm_model = load_segmentation_model()
>>> segm_map, key_maps = segment_subcortical(aligned_scan_prep, segm_model)
>>> assert segm_map.shape == (1, 160, 160, 160)