Self-supervision Based Dual-Transformation Learning for Stain Normalization, Classification andSegmentation

Shiv Gehlot, Anubha Gupta

Abstract: Stain color variation s across images are common in the medical imaging domain. However, such variations among the training and test datasets may lead to unsatisfactory performance on the latter in any desired task. This paper proposes a novel coupled-network composed of two U-Net type architectures that utilize self-supervised learning. The first subnetwork (N1) learns an identity transformation, while the second (N2) learns a transformation to perform stain normalization. We also introduce classification heads in the subnetworks, trained along with the stain normalization task. To the best of our knowledge, the proposed coupling framework, where the information from the encoders of both the subnetworks is utilized by the decoders of both subnetworks as well as trained in a coupled fashion, is introduced in this domain for the first time. Interestingly, the coupling of N1 (for identity transformation) and N2 (for stain normalization) helps N2 learn the stain normalization task while being cognizant of the features essential to reconstruct images. Similarly, N1 learns to extract relevant features for reconstruction invariant to stain color variations due to its coupling with N2. Thus, the two subnetworks help each other, leading to improved performance on the subsequent task of classification. Further, it is shown that the proposed architecture can also be used for segmentation, making it applicable for all three applications: stain normalization, classification, and segmentation. Experiments are carried out on four datasets to show the efficacy of the proposed architecture.