Domain-invariant Representation Learning via Segment Anything Model for Blood Cell Classification

TL;DR

This work tackles domain shifts in blood cell classification by introducing DoRL, a framework that merges a LoRA-tuned Segment Anything Model (LoRA-SAM) for segmentation with a cross-domain autoencoder (CAE) to learn domain-invariant features. The CAE employs MAE-inspired masking, a reconstruction loss $\mathcal{L}_{rec}$, a discrepancy loss $\mathcal{L}_{mmd}$, and a structural similarity loss $\mathcal{L}_{ssim}$, producing features $s_i$ that are robust across datasets. DoRL’s representations are evaluated with five classifiers (RF, SVM, LR, ANN, XGBoost), achieving state-of-the-art cross-domain performance on Matek-19, Acevedo-20, and SYSU3H, notably with SVM (linear) reaching $60.36\%$ average accuracy. The approach also demonstrates artifact suppression and clear clustering in latent space via visualizations, underscoring its potential for real-world clinical deployment. The code is available at the provided URL, highlighting the method’s practicality and reproducibility.

Abstract

Accurate classification of blood cells is of vital significance in the diagnosis of hematological disorders. However, in real-world scenarios, domain shifts caused by the variability in laboratory procedures and settings, result in a rapid deterioration of the model's generalization performance. To address this issue, we propose a novel framework of domain-invariant representation learning (DoRL) via segment anything model (SAM) for blood cell classification. The DoRL comprises two main components: a LoRA-based SAM (LoRA-SAM) and a cross-domain autoencoder (CAE). The advantage of DoRL is that it can extract domain-invariant representations from various blood cell datasets in an unsupervised manner. Specifically, we first leverage the large-scale foundation model of SAM, fine-tuned with LoRA, to learn general image embeddings and segment blood cells. Additionally, we introduce CAE to learn domain-invariant representations across different-domain datasets while mitigating images' artifacts. To validate the effectiveness of domain-invariant representations, we employ five widely used machine learning classifiers to construct blood cell classification models. Experimental results on two public blood cell datasets and a private real dataset demonstrate that our proposed DoRL achieves a new state-of-the-art cross-domain performance, surpassing existing methods by a significant margin. The source code can be available at the URL (https://github.com/AnoK3111/DoRL).

Figures (7)

• Figure 1: The grayscale distribution of different blood cell image datasets. The blood cell image datasets often exhibit domain shift issues, resulting in the deterioration of model generalization performance. For instance, when trained on the Matek-19 dataset, ResNext can achieve impressive performance on the source dataset. However, its performance deteriorates significantly when applied to unseen datasets (Acevedo-20 and SYSU3H).
• Figure 2: The pipeline of our proposed blood cell classification method. It comprises two primary components: domain-invariant representation learning and blood cell classification. In the stage of domain-invariant representation learning, DoRL is employed to extract domain-invariant features from blood cell images. To evaluate the effectiveness of these domain-invariant features, five widely used machine learning classifiers are employed to construct blood cell classification models.
• Figure 3: The architecture of our proposed DoRL. It consists of two key components: the LoRA-SAM and a CAE. The LoRA-SAM component is tasked with learning image embeddings from blood cell images and performing segmentation on these images. Subsequently, the segmented blood cell images and their corresponding embeddings are input into the CAE, which extracts domain-invariant features from various blood cell datasets while eliminating image artifacts.
• Figure 4: The statistic distribution of blood cell image datasets.
• Figure 5: The visual presentation of original, segmented, and reconstructed images. Notably, the second row is the segmented images generated by LoRA-SAM and the third row is the images reconstructed by CAE. It is observed that our proposed DoRL can effectively eliminate artifacts and accurately reconstruct blood cells.