Leveraging Self-attention Mechanism in Vision Transformers for Unsupervised Segmentation of Optical Coherence Microscopy White Matter Images

A new microscope has been created to capture detailed images of the brain using a technology called optical coherence microscopy (OCM). However, there is still much to discover and understand about this valuable data. In this paper, we focus on the important task of segmenting the white matter in these high-resolution OCM images. A closed-up accurate segmentation of white matter tracts has the potential to enhance our knowledge of brain connections. In this paper, we propose an unsupervised segmentation approach that leverages the self-attention mechanism of Vision Transformers (ViT). Our approach uses the output attention weights from a ViT pre-trained with Masked Image Modeling (MIM) to generate binary segmentations that we use as Pseudo-Ground-Truth (PGT) to train an additional segmentation model. Our method achieved superior performance when compared with classical unsupervised computer vision methods and common unsupervised deep learning architectures designed for natural images. Additionally, we compared our results with those of a supervised U-Net model trained on different numbers of labels and a semi-supervised approach where we selected the best-performing model based on labeled data. Our model achieved comparable results to the U-Net model trained on 30% of the labeled data. Furthermore, through fine-tuning, our model demonstrated an improvement of 3% over the supervised U-Nets. The code and data are available on GitHub repository https://github.com/linum-uqam/ViT-OCM-WMSegmentation .