Gradient-Rebalanced Uncertainty Minimization for Cross-Site Adaptation of Medical Image Segmentation

Jiaming Li, Chaowei Fang, Guanbin Li

Abstract: Automatically adapting image segmentation across data sites benefits to reduce the data annotation burden in medical image analysis. Due to variations in image collection procedures, there usually exists moderate domain gap between medical image datasets from different sites. Increasing the prediction certainty is beneficial for gradually reducing the category-wise domain shift. However, uncertainty minimization naturally leads to bias towards major classes since the target object usually occupies a small portion of pixels in the input image. In this paper, we propose a gradient-rebalanced uncertainty minimization scheme which is capable of eliminating the learning bias. First, the foreground pixels and background pixels are reweighted according to the total gradient amplitude of every class. Furthermore, we devise a feature-level adaptation scheme to reduce the overall domain gap between source and target datasets, based on feature norm regularization and adversarial learning. Experiments on CT pancreas segmentation and MRI prostate segmentation validate that, our method outperforms existing cross-site adaptation algorithms by around 3% on the DICE similarity coefficient.