cellmil.models.segmentation.hovernet

Copyright 2021, Dana-Farber Cancer Institute and Weill Cornell Medicine License: GNU GPL 2.0

Functions

compute_hv_map(mask)	Preprocessing step for HoVer-Net architecture.
extract_nuclei_info(pred_inst[, pred_type])	Extract centroids, type, and probability of the type for each cell in the nuclei mask.
group_centroids_by_type(cell_dict, ...)	Group centroids by cell type for cells above a certain probability threshold.
loss_hovernet(outputs, ground_truth[, n_classes])	Compute loss for HoVer-Net.
post_process_batch_hovernet(outputs, n_classes)	Post-process HoVer-Net outputs to get a final predicted mask.
remove_small_objs(array_in, min_size)	Removes small foreground regions from binary array, leaving only the contiguous regions which are above the size threshold.

Classes

HoVerNet([n_classes])

Model for simultaneous segmentation and classification based on HoVer-Net.

class cellmil.models.segmentation.hovernet._BatchNormRelu(n_channels: int)

Bases: Module

BatchNorm + Relu layer

__init__(n_channels: int)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class cellmil.models.segmentation.hovernet._HoVerNetResidualUnit(input_channels: int, output_channels: int, stride: int)

Bases: Module

Residual unit. See: Fig. 2(a) from Graham et al. 2019 HoVer-Net paper. This unit is not preactivated! That’s handled when assembling units into blocks. output_channels corresponds to m in the figure

__init__(input_channels: int, output_channels: int, stride: int)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

cellmil.models.segmentation.hovernet._make_HoVerNet_residual_block(input_channels: int, output_channels: int, stride: int, n_units: int)

Stack multiple residual units into a block. output_channels is given as m in Fig. 2 from Graham et al. 2019 paper

class cellmil.models.segmentation.hovernet._HoVerNetEncoder

Bases: Module

Encoder for HoVer-Net. 7x7 conv, then four residual blocks, then 1x1 conv. BatchNormRelu after first convolution, based on code from authors, see: (https://github.com/vqdang/hover_net/blob/5d1560315a3de8e7d4c8122b97b1fe9b9513910b/src/model/graph.py#L67)

Return a list of the outputs from each residual block, for later skip connections

__init__()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: Tensor) → List[Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class cellmil.models.segmentation.hovernet._HoVerNetDenseUnit(input_channels: int)

Bases: Module

Dense unit. See: Fig. 2(b) from Graham et al. 2019 HoVer-Net paper.

__init__(input_channels: int)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

cellmil.models.segmentation.hovernet._make_HoVerNet_dense_block(input_channels: int, n_units: int)

Stack multiple dense units into a block.

class cellmil.models.segmentation.hovernet._HoverNetDecoder

Bases: Module

One of the three identical decoder branches.

__init__()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: List[Tensor]) → Tensor

Inputs should be a list of the outputs from each residual block, so that we can use skip connections

class cellmil.models.segmentation.hovernet.HoVerNet(n_classes: int | None = 6)

Bases: Module

Model for simultaneous segmentation and classification based on HoVer-Net. Can also be used for segmentation only, if class labels are not supplied. Each branch returns logits.

Parameters:

n_classes (int) – Number of classes for classification task. If None then the classification branch is not used.

References

Graham, S., Vu, Q.D., Raza, S.E.A., Azam, A., Tsang, Y.W., Kwak, J.T. and Rajpoot, N., 2019. Hover-Net: Simultaneous segmentation and classification of nuclei in multi-tissue histology images. Medical Image Analysis, 58, p.101563.

__init__(n_classes: int | None = 6)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs: Tensor) → dict[str, torch.Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

calculate_instance_map(predictions: dict[str, torch.Tensor], magnification: int | float = 40) → Tuple[Tensor, list[dict[numpy.int32, dict[str, Any]]]]

Calculate Instance Map from network predictions (after Softmax output)

Parameters:

• predictions (dict) – Dictionary with the following required keys: * nuclei_binary_map: Binary Nucleus Predictions. Shape: (batch_size, H, W, 2) * nuclei_type_map: Type prediction of nuclei. Shape: (batch_size, H, W, 6) * hv_map: Horizontal-Vertical nuclei mapping. Shape: (batch_size, H, W, 2)

• magnification (Literal[20, 40], optional) – Which magnification the data has. Defaults to 40.

Returns:

• torch.Tensor: Instance map. Each Instance has own integer. Shape: (batch_size, H, W)

• List of dictionaries. Each List entry is one image. Each dict contains another dict for each detected nucleus.

  For each nucleus, the following information are returned: “bbox”, “centroid”, “contour”, “type_prob”, “type”

Return type:

Tuple[torch.Tensor, List[dict]]

cellmil.models.segmentation.hovernet._convert_multiclass_mask_to_binary(mask: Tensor) → Tensor

Input mask of shape (B, n_classes, H, W) is converted to a mask of shape (B, 1, H, W). The last channel is assumed to be background, so the binary mask is computed by taking its inverse.

cellmil.models.segmentation.hovernet._dice_loss_np_head(np_out: Tensor, true_mask: Tensor, epsilon: float = 0.001)

Dice loss term for nuclear pixel branch. This will compute dice loss for the entire batch (not the same as computing dice loss for each image and then averaging!)

Parameters:

• np_out – logit outputs of np branch. Tensor of shape (B, 2, H, W)

• true_mask – True mask. Tensor of shape (B, n_classes, H, W)

• epsilon (float) – Epsilon passed to dice_loss()

cellmil.models.segmentation.hovernet._dice_loss_nc_head(nc_out: Tensor, true_mask: Tensor, epsilon: float = 0.001)

Dice loss term for nuclear classification branch. Computes dice loss for each channel, and sums up. This will compute dice loss for the entire batch (not the same as computing dice loss for each image and then averaging!)

Parameters:

• nc_out – logit outputs of nc branch. Tensor of shape (B, n_classes, H, W)

• true_mask – True mask. Tensor of shape (B, n_classes, H, W)

• epsilon (float) – Epsilon passed to dice_loss()

cellmil.models.segmentation.hovernet._ce_loss_nc_head(nc_out: Tensor, true_mask: Tensor)

Cross-entropy loss term for nc branch. :param nc_out: logit outputs of nc branch. Tensor of shape (B, n_classes, H, W) :param true_mask: True mask. Tensor of shape (B, n_classes, H, W)

cellmil.models.segmentation.hovernet._ce_loss_np_head(np_out: Tensor, true_mask: Tensor)

Cross-entropy loss term for np branch. :param np_out: logit outputs of np branch. Tensor of shape (B, 2, H, W) :param true_mask: True mask. Tensor of shape (B, n_classes, H, W)

cellmil.models.segmentation.hovernet.compute_hv_map(mask: ndarray[Any, Any]) → ndarray[Any, Any]

Preprocessing step for HoVer-Net architecture. Compute center of mass for each nucleus, then compute distance of each nuclear pixel to its corresponding center of mass. Nuclear pixel distances are normalized to (-1, 1). Background pixels are left as 0. Operates on a single mask. Can be used in Dataset object to make Dataloader compatible with HoVer-Net.

Based on https://github.com/vqdang/hover_net/blob/195ed9b6cc67b12f908285492796fb5c6c15a000/src/loader/augs.py#L192

Parameters:

mask (np.ndarray) – Mask indicating individual nuclei. Array of shape (H, W), where each pixel is in {0, …, n} with 0 indicating background pixels and {1, …, n} indicating n unique nuclei.

Returns:

array of hv maps of shape (2, H, W). First channel corresponds to horizontal and second vertical.

Return type:

np.ndarray

cellmil.models.segmentation.hovernet._get_gradient_hv(hv_batch: Tensor, kernel_size: int = 5)

Calculate the horizontal partial differentiation for horizontal channel and the vertical partial differentiation for vertical channel. The partial differentiation is approximated by calculating the central differnce which is obtained by using Sobel kernel of size 5x5. The boundary is zero-padded when channel is convolved with the Sobel kernel.

Parameters:

• hv_batch – tensor of shape (B, 2, H, W). Channel index 0 for horizonal maps and 1 for vertical maps. These maps are distance from each nuclear pixel to center of mass of corresponding nucleus.

• kernel_size (int) – width of kernel to use for gradient approximation.

Returns:

Tuple of (h_grad, v_grad) where each is a Tensor giving horizontal and vertical gradients respectively

cellmil.models.segmentation.hovernet._loss_hv_grad(hv_out: Tensor, true_hv: Tensor, nucleus_pixel_mask: Tensor)

Equation 3 from HoVer-Net paper for calculating loss for HV predictions. Mask is used to compute the hv loss ONLY for nuclear pixels

Parameters:

• hv_out – Ouput of hv branch. Tensor of shape (B, 2, H, W)

• true_hv – Ground truth hv maps. Tensor of shape (B, 2, H, W)

• nucleus_pixel_mask – Boolean mask indicating nuclear pixels. Tensor of shape (B, H, W)

cellmil.models.segmentation.hovernet._loss_hv_mse(hv_out: Tensor, true_hv: Tensor)

Equation 2 from HoVer-Net paper for calculating loss for HV predictions.

Parameters:

• hv_out – Ouput of hv branch. Tensor of shape (B, 2, H, W)

• true_hv – Ground truth hv maps. Tensor of shape (B, 2, H, W)

cellmil.models.segmentation.hovernet.loss_hovernet(outputs: list[torch.Tensor], ground_truth: list[torch.Tensor], n_classes: Optional[int] = None)

Compute loss for HoVer-Net. Equation (1) in Graham et al.

Parameters:

• outputs –

  Output of HoVer-Net. Should be a list of [np, hv] if n_classes is None, or a list of [np, hv, nc] if n_classes is not None. Shapes of each should be:

  • np: (B, 2, H, W)

  • hv: (B, 2, H, W)

  • nc: (B, n_classes, H, W)

• ground_truth – True labels. Should be a list of [mask, hv], where mask is a Tensor of shape (B, 1, H, W) if n_classes is None or (B, n_classes, H, W) if n_classes is not None. hv is a tensor of precomputed horizontal and vertical distances of nuclear pixels to their corresponding centers of mass, and is of shape (B, 2, H, W).

• n_classes (int) – Number of classes for classification task. If None then the classification branch is not used.

References

Graham, S., Vu, Q.D., Raza, S.E.A., Azam, A., Tsang, Y.W., Kwak, J.T. and Rajpoot, N., 2019. Hover-Net: Simultaneous segmentation and classification of nuclei in multi-tissue histology images. Medical Image Analysis, 58, p.101563.

cellmil.models.segmentation.hovernet.remove_small_objs(array_in: ndarray[Any, Any], min_size: int)

Removes small foreground regions from binary array, leaving only the contiguous regions which are above the size threshold. Pixels in regions below the size threshold are zeroed out.

Parameters:

• array_in (np.ndarray) – Input array. Must be binary array with dtype=np.uint8.

• min_size (int) – Minimum size of each region.

Returns:

Array of labels for regions above the threshold. Each separate contiguous region is labelled with

a different integer from 1 to n, where n is the number of total distinct contiguous regions

Return type:

np.ndarray

cellmil.models.segmentation.hovernet._post_process_single_hovernet(np_out: Tensor, hv_out: Tensor, small_obj_size_thresh: int = 10, kernel_size: int = 21, h: float = 0.5, k: float = 0.5) → ndarray[Any, Any]

Combine predictions of np channel and hv channel to create final predictions. Works by creating energy landscape from gradients, and the applying watershed segmentation. This function works on a single image and is wrapped in post_process_batch_hovernet() to apply across a batch. See: Section B of HoVer-Net article and https://github.com/vqdang/hover_net/blob/14c5996fa61ede4691e87905775e8f4243da6a62/models/hovernet/post_proc.py#L27

Parameters:

• np_out (torch.Tensor) – Output of NP branch. Tensor of shape (2, H, W) of logit predictions for binary classification

• hv_out (torch.Tensor) – Output of HV branch. Tensor of shape (2, H, W) of predictions for horizontal/vertical maps

• small_obj_size_thresh (int) – Minimum number of pixels in regions. Defaults to 10.

• kernel_size (int) – Width of Sobel kernel used to compute horizontal and vertical gradients.

• h (float) – hyperparameter for thresholding nucleus probabilities. Defaults to 0.5.

• k (float) – hyperparameter for thresholding energy landscape to create markers for watershed segmentation. Defaults to 0.5.

cellmil.models.segmentation.hovernet.post_process_batch_hovernet(outputs: list[torch.Tensor], n_classes: int | None, small_obj_size_thresh: int = 10, kernel_size: int = 21, h: float = 0.5, k: float = 0.5, return_nc_out_preds: bool = False) → Union[ndarray[Any, Any], Tuple[ndarray[Any, Any], ndarray[Any, Any]], Tuple[ndarray[Any, Any], ndarray[Any, Any], ndarray[Any, Any]]]

Post-process HoVer-Net outputs to get a final predicted mask. See: Section B of HoVer-Net article and https://github.com/vqdang/hover_net/blob/14c5996fa61ede4691e87905775e8f4243da6a62/models/hovernet/post_proc.py#L27

Parameters:

• outputs (list) –

  Outputs of HoVer-Net model. List of [np_out, hv_out], or [np_out, hv_out, nc_out] depending on whether model is predicting classification or not.

  • np_out is a Tensor of shape (B, 2, H, W) of logit predictions for binary classification

  • hv_out is a Tensor of shape (B, 2, H, W) of predictions for horizontal/vertical maps

  • nc_out is a Tensor of shape (B, n_classes, H, W) of logits for classification

• n_classes (int) – Number of classes for classification task. If None then only segmentation is performed.

• small_obj_size_thresh (int) – Minimum number of pixels in regions. Defaults to 10.

• kernel_size (int) – Width of Sobel kernel used to compute horizontal and vertical gradients.

• h (float) – hyperparameter for thresholding nucleus probabilities. Defaults to 0.5.

• k (float) – hyperparameter for thresholding energy landscape to create markers for watershed segmentation. Defaults to 0.5.

Returns:

If n_classes is None, returns det_out. In classification setting, returns (det_out, class_out).

• det_out is np.ndarray of shape (B, H, W)

• class_out is np.ndarray of shape (B, n_classes, H, W)

Each pixel is labelled from 0 to n, where n is the number of individual nuclei detected. 0 pixels indicate background. Pixel values i indicate that the pixel belongs to the ith nucleus.

Return type:

np.ndarray

Modified previous method to output nc_out_preds.

cellmil.models.segmentation.hovernet.extract_nuclei_info(pred_inst: ndarray[Any, Any], pred_type: Optional[ndarray[Any, Any]] = None) → dict[int, dict[str, Any]]

Extract centroids, type, and probability of the type for each cell in the nuclei mask.

This function is adapted from the Hover-Net post-processing implementation. Original source: https://github.com/vqdang/hover_net/blob/14c5996fa61ede4691e87905775e8f4243da6a62/models/hovernet/post_proc.py#L94

Parameters:

• pred_inst (np.ndarray) – Instance mask of nuclei.

• pred_type (np.ndarray) – Optional type mask for nuclei.

Returns:

Information about each cell, including centroid, type, and probability.

Return type:

dict

cellmil.models.segmentation.hovernet.group_centroids_by_type(cell_dict: dict[str, Any], prob_threshold: float)

Group centroids by cell type for cells above a certain probability threshold.

Parameters:

• cell_dict (dict) – Dictionary containing cell information.

• prob_threshold (float) – Minimum probability threshold for a cell to be considered.

Returns:

A dictionary with cell types as keys and lists of centroids as values.

Return type:

dict