CateNorm: Categorical Normalization for Robust Medical Image Segmentation

Abstract

Batch normalization (BN) uniformly shifts and scales the activations based on the statistics of a batch of images. However, the intensity distribution of the background pixels often dominates the BN statistics because the background accounts for a large proportion of the entire image. This paper focuses on enhancing BN with the intensity distribution of foreground pixels, the one that really matters for image segmentation. We propose a new normalization strategy, named categorical normalization (CateNorm), to normalize the activations according to categorical statistics. The categorical statistics are obtained by dynamically modulating specific regions in an image that belong to the foreground. CateNorm demonstrates both precise and robust segmentation results across five public datasets obtained from different domains, covering complex and variable data distributions. It is attributable to the ability of CateNorm to capture domain-invariant information from multiple domains (institutions) of medical data.

Code is available at https://github.com/lambert-x/CateNorm.

Appendices

A Details of Aligning Input Distribution Algorithm

Assume that we have N source domains $S_1, S_2, S_3, ..., S_N$, with $M_1, M_2, M_3, ..., M_N$ examples respectively, where the i-th domain source domain $S_i$ consists of an image set $\{\textbf{x}_{i,j}\in \mathbb {R}^{D_{i,j}}\}_{j=1,...,M_i}$ as well as their associated annotations. Our goal is to align the image distributions of these source domains with the target domain T based on the class-wise (region-wise) statistics. The algorithm can be illustrated as the following steps:

Step 1: Calculate class-wise statistics of each case

Firstly, we calculate the mean and standard deviation of each case in both the source domain and the target domain.

$$\begin{aligned} \mu _{i,j}^{c}&= \frac{\sum _{k=1}^{|D^c_{i,j}|}\textbf{x}^c_{i,j,k}}{|D^c_{i,j}|} , \end{aligned}$$

(6)

$$\begin{aligned} \sigma _{i,j}^{c}&= \sqrt{\frac{1}{ |D^{c}_{i,j}|} \sum _{k=1}^{|D^{c}_{i,j}|}(\textbf{x}^c_{i,j,k} - \mu _{i,j}^{c})^{2}}, \end{aligned}$$

(7)

where $\textbf{x}_{i,j}^{c}$ denotes the pixels which belong to the c-th class (region) in image $\textbf{x}_{i,j}$, with the number of pixels denoted as $|D^c_{i,j}|$. As a special case, $i=T$ indicates the target domain.

Step 2: Estimate aligned (new) class-wise statistics

Next, we calculate the mean of the statistics over all examples obtained in each domain as follows:

$$\begin{aligned} \bar{\mu }_{i}^{c}&=\frac{\sum _{j=1}^{M_i}\mu _{i,j}^{c}}{M_i}, \end{aligned}$$

(8)

$$\begin{aligned} \bar{\sigma }_{i}^{c}&=\frac{\sum _{j=1}^{M_i}\sigma _{i,j}^{c}}{M_i}. \end{aligned}$$

(9)

Based on the $\bar{\mu }_{i}^{c}$, we now estimate the new class-wise mean $\tilde{\mu }_{i,j}$ for each case of the source domain $S_i$ as follows:

$$\begin{aligned} \begin{aligned} \tilde{\mu }_{i,j}^{c}&=\frac{\mu _{i,j}^{c} - \bar{\mu }_{i}^{c}}{\sqrt{\frac{\sum _{j=1}^{M_i}(\mu _{i,j}^{c} - \bar{\mu }_{i}^{c})^2}{M_i}}} \cdot \sqrt{\frac{\sum _{j=1}^{M_T}(\mu _{T,j}^{c} - \bar{\mu }_{T}^{c})^2}{M_T}} + \bar{\mu }_{T}^{c}, \\ \end{aligned} \end{aligned}$$

(10)

where $M_T$ denotes the number of cases in the target domain T. Similarly, the new standard deviation $\tilde{\sigma }_{i,j}$ can be computed by:

$$\begin{aligned} \begin{aligned} \tilde{\sigma }_{i,j}^{c}&=\frac{\sigma _{i,j}^{c} - \bar{\sigma }_{i}^{c}}{\sqrt{\frac{\sum _{j=1}^{M_i}(\sigma _{i,j}^{c} - \bar{\sigma }_{i}^{c})^2}{M_i}}} \cdot \sqrt{\frac{\sum _{j=1}^{M_T}(\sigma _{T,j}^{c} - \bar{\sigma }_{T}^{c})^2}{M_T}} + \bar{\sigma }_{T}^{c}. \\ \end{aligned} \end{aligned}$$

(11)

Step 3: Align each case with the estimated statistics

Based on the computed new mean and standard deviation $\tilde{\mu }_{i,j}$, $\tilde{\sigma }_{i,j}$, the aligned image $\tilde{\textbf{x}}_{i,j}$ can be computed as:

$$\begin{aligned} \tilde{\textbf{x}}_{i,j}^{c}&=\frac{\textbf{x}_{i,j}^{c} - \mu _{i,j}^{c}}{\sigma _{i,j}^{c}} \cdot \tilde{\sigma }_{i,j}^{c} + \tilde{\mu }_{i,j}^{c}. \end{aligned}$$

(12)

B Implementation Details

(See Tables 4, 5).

Table 4. Data preprocessing.

Table 5. Experimental setting.

C Details of the prostate datasets

(See Tables 6).

Table 6. Details of the 3 prostate segmentation datasets.

D Training procedure of CateNorm

E Average Surface Distance (ASD) Comparison

The detailed average surface distance results of both prostate segmentation and abdominal segmentation tasks can be found in Tables 7 and 8. The proposed CateNorm achieves the lowest average ASD on both tasks, even under the more challenging multi-domain setting (Tables 9, 10, 12).

Table 7. ASD comparison on the abdominal datasets under the multi-domain setting (in mm). Compared with the baseline and other competitive methods, the proposed CateNorm achieves the lowest average ASD.

Table 8. ASD comparison on prostate segmentation datasets under the multi-domain setting (in mm). Compared with the baseline and other competitive methods, the proposed CateNorm achieves the lowest average ASD.

Fig. 3.

Performance gain under partial annotation. We compare our method to the baseline with fewer annotated classes (i.e., 3/5). We can see that by partitioning the images into different number of regions, CateNorm consistently achieves better results than BN for all tested organs. This suggests that our algorithm is not sensitive to the number of regions.

Fig. 4.

Qualitative results comparison. We compare our baseline and the other SOTA method under the multi-domain setting on prostate segmentation and abdominal multi-organ segmentation. Results in the first three rows clearly show that our method outperforms others as their results are cracked and incomplete with these unapparent prostate boundaries. And the results in the last two rows show our methods could better suppress inconsistent class information inside a close segmented area (e.g., reducing false positives inside the stomach) and predict hard organs like the pancreas more accurately by incorporating general and categorical statistics.

Table 9. Comparison on the multi-organ segmentation dataset (BTCV) with single-domain setting (Dice Score in %).

Table 10. Organ-wise results on the multi-organ segmentation datasets under the multi-domain setting (Dice Score in %).

Table 11. CateNorm is compatible to other segmentation models. This table compares performance on multi-domain multi-organ and prostate segmentation with DeepLabv3+ [5] architecture. Our CateNorm consistently outperforms BN.

Table 12. CateNorm is not sensitive to the warmup length. This table reports average accuracy (%) of our CateNorm under deteriorated pretrained models with fewer pretraining iterations. We reduce the warmup iterations to 450, 1440, and 9000 for multi-domain prostate segmentation experiments, to investigate how our CateNorm performs when warmuped with less iterations.

Fig. 5.

Set CateNorm block(s) early. This table compares performance with single CateNorm block set in different positions. Adding the CateNorm to the encoder (block index 1–5) always yields better performance than adding to the decoder (block index 6–10). In general, the performance decreases as the block index increases. We believe that it is because the earlier layers in the encoder extract lower-level features that are less discriminative than the decoder features.

Fig. 6.

CateNorm does normalize with semantic information. This figure visualizes the learned $\gamma ^{CateNorm }$ (1st row) and $\beta ^{CateNorm }$ (2nd row) of a CateNorm layer in a CateNorm block on different channels of the intermediate CateNorm layer during the second forward. With prior class information as guidance, CateNorm can modulate spatially-adaptive parameters. Such spatial-wise modulation can be complementary to the channel-wise modulation accomplished by BN, and derives more discriminative features that benefit segmentation.