Addressing Domain Shift in CNN-based Image Segmentation: From Improving Robustness to Unsupervised and Active Learning based Domain Adaptation

Abstract

Pixel-wise labeling of images is a common task in applications nowadays as it allows humans and machines to make sense of the content in an image and informs their decisions about subsequent steps to be taken. Examples of such applications include scene understanding for autonomous driving, weeding in agricultural fields, and interventional therapies for brain tumors. It is however challenging to build machine learning models that perform well on data exhibiting different characteristics to what they have seen during training, or put differently, that they are robust to domain shift. <br /> We address this problem using different approaches. A starting point is to increase the robustness of a segmentation model by encouraging it to focus more on the shape content of images rather than textural features. We accomplish this by using TV augmentation to smooth images while emphasizing object boundaries. We compare our work to other augmentation techniques and show the benefit of our approach. Another aspect we look into is design decisions employed by a popular open-source framework that is widely used for segmentation of medical images. Concretely, we investigate the effect of optimizers and the number of training epochs on domain generalization, and show through our experiments how the performance on new domains can be improved. <br /> If labeled images from the target domain can be acquired, then these can be used to adapt models to the unseen domains. We devise smart sampling strategies to select which data samples should be annotated for efficient active learning. We first generate pseudo-labels and choose samples based on the loss and gradients of the network. Finally, if only unlabeled images are available, we use self-supervision to leverage this data and adapt the model to the target domain. We add a second branch to the trained model and drive the optimization of the model by comparing two sets of segmentation maps on target data. Our use of probabilistic maps and limited supervision reduces the risk of propagating incorrect signals throughout the network while allowing the model to adapt itself to target data.