Automatic Recognition of White Blood Cell Images with Memory Efficient Superpixel Metric GNN: SMGNN

1. Methodolgy of Superpixel Metric

Deep learning methods for medical image processing have predominantly concentrated on discerning the local context, which refers to inter-pixel dependencies within individual images [56]. However, there’s a missed opportunity: capturing the “global” context that exists between training samples. While pixel-level contrast or metric learning provides a way to bridge this gap, the sheer computational and memory overheads—due to contrast or metric computations spanning every pixel pair—render them less feasible. We propose an innovative efficient superpixel-level metric learning on metric loss, which not only captures the desired global context but does so while drastically cutting computational and memory costs.

1.1. Compression Ratio on Image Data

The utility of superpixel methods in image data preprocessing is well-acknowledged, particularly for their ability to condense data and reduce computational demands. Consequently, this study capitalizes on these advantages, transforming the image data into a more compact graph representation. Within this framework, each superpixel evolves into a graph node. The interconnectedness of these nodes—whether driven by spatial positioning or feature similarity—determines the graph’s topology. These node features aren’t rigid; their definition can range from basic 5-dimensional attributes encompassing color (3 dimensions) and location (2 dimensions) to more intricate data points like histograms, positional variance, and variations in pixel values.

Given that the adjacency matrix exhibits sparse characteristics, it’s predominantly the node features that dictate memory consumption. Opting for the more straightforward features facilitates significant data compression. Specifically, for 3-channel RGB images, we’ve discerned a compression ratio roughly represented as $c=\frac{\mathrm{graph}\mathrm{data}}{\mathrm{image}\mathrm{data}}\approx \frac{5K}{3n}$. Importantly, this efficient compression does not compromise on quality. Our subsequent numerical experiments demonstrate that this ratio is concomitant with optimal segmentation outcomes.

1.2. Quality of Superpixel and Reconstruction Score

The segmentation task relies on the quality of the generated superpixels, and good superpixel results coherent with the boundary of labeled images. Suppose $\mathbf{X}\in {\mathbb{R}}^{H\times W\times 3}$ is the input image. Let $\mathcal{V}$ be the set of all superpixels, $N=H\times W$ be the number of pixels, $K=\left|\mathcal{V}\right|$ be the number of superpixels, and $\mathbf{Q}\in {\mathbb{R}}^{N\times K}$ be the association matrix between pixels and superpixels, then we have

$${\mathbf{Q}}_{i,j}=\left\{\begin{array}{cc}1,\hfill & \mathrm{if}{\widehat{\mathbf{X}}}_i\in {\mathcal{V}}_j\hfill \\ 0,\hfill & \mathrm{otherwise}\hfill \end{array},\widehat{\mathbf{X}}=Flatten\left(\mathbf{X}\right)\in {\mathbb{R}}^{N\times 3},\right.$$

where ${\mathcal{V}}_j$ is the jth superpixel, Flatten operation converts a two-dimensional image matrix into a one-dimensional vector. The role of the association matrix build the bridge between image space and graph space. For computation convenience, we define column normalized association matrix as

$${\overline{\mathbf{Q}}}_{i,j}=\left\{\begin{array}{cc}1/\left[{\mathcal{S}}_j\right],\hfill & \mathrm{if}{\widehat{\mathbf{X}}}_i\in {\mathcal{V}}_j\hfill \\ 0,\hfill & \mathrm{otherwise}\hfill \end{array},\left[{\mathcal{S}}_j\right]=\sum _{x\in {\mathcal{V}}_j}1.\right.$$

Let $\mathbf{Y}\in {\mathbb{R}}^N$ be the label of the image pixels and $\mathcal{Y}\in {\mathbb{R}}^K$ be superpixel metric or metric score of the graph data. Using the pixel label, we can formulate metric score as

$${\mathcal{Y}}_j=\frac{{\sum }_{x\in {\mathcal{V}}_j}{\mathbf{Y}}_x}{{\sum }_{x\in {\mathcal{V}}_j}1},$$

(1)

for the $jth$ superpixel. $\mathcal{Y}$ can also be efficiently computed with column normalized association matrix, i.e., $\mathcal{Y}={\overline{\mathbf{Q}}}^T\mathbf{Y}$. We can back-project the supepixel label to image space by association matrix, i.e., $\overline{\mathbf{Y}}=\mathbf{Q}\mathcal{Y}$. We define the IoU reconstruction score to evaluate the quality, which reads

$$Io{U}_r=\frac{1}{c}\sum _{i=0}^c\frac{\left\{x\right|{\overline{\mathbf{Y}}}_x=i\}\cap \left\{x\right|{\mathbf{Y}}_x=i\}}{\left\{x\right|{\overline{\mathbf{Y}}}_x=i\}\cup \left\{x\right|{\mathbf{Y}}_x=i\}},$$

(2)

where $\left\{x\right|{\mathbf{Y}}_x=i\}$ denotes the set of pixels whose class label equal i, c is the number of classes.

1.3. Lightweight Graph Neural Networks for Superpixel Embedding

Graph neural networks have the advantage of being lightweight compared to other deep learning models that often require deeper and larger networks to achieve higher performance. GNNs leverage relational information and utilize shallow layers to achieve satisfactory results.

Given an undirected attributed graph $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathcal{S})$, $\mathcal{G}$ consists of a nonempty finite set of $K=\left|\mathcal{V}\right|$ nodes $\mathcal{V}$ and a set of edges $\mathcal{E}$ between node pairs. Denote $\mathcal{A}\in {\mathbb{R}}^{K\times K}$ the graph adjacency matrix and $\mathcal{S}\in {\mathbb{R}}^{K\times d}$ the node attributes. A graph convolution learns a matrix representation $\mathcal{H}$ that embeds the structure $\mathcal{A}$ and feature matrix $\mathcal{S}={\left\{{\mathcal{S}}_j\right\}}_{j=1}^N$ with ${\mathcal{S}}_j$ for node j. Most graph convolutions follow the message passing [57] update scheme, which finds a central node’s smooth representation by aggregating its 1-hop neighbor information. At layer ℓ, the propagation for the ith node reads

$${\mathcal{H}}_i^{\ell }=\gamma \left({\mathcal{H}}_i^{\ell -1},{\square }_{j\in \mathcal{N}\left(i\right)}\varphi \left({\mathcal{H}}_i^{\ell -1},{\mathcal{H}}_j^{\ell -1},{\mathcal{A}}_{ij}\right)\right),{\mathcal{H}}^0=\mathcal{S}$$

(3)

where $\square (·)$ is a differentiable and permutation invariant aggregation function, such as summation, average, or maximization. The set $\mathcal{N}\left(i\right)$ includes ${\mathcal{V}}_i$ and its 1-hop neighbors. Both $\gamma (·)$ and $\varphi (·)$ are differentiable aggregation functions, such as MultiLayer Perceptrons (MLPs).

In our approach, we construct graph data using superpixels and their predefined relationships. GNNs learn features from the graph space to enhance the segmentation capability. We employ the Graph Isomorphism Network (GIN) [26] as the backbone graph representation network. At each layer ℓ, the GIN model updates the ith node representation as follows:

$${\mathcal{H}}_i^{(\ell )}={\mathrm{MLP}}^{(\ell )}\left(\left(1+{\mathbf{w}}^{(\ell )}\right)·{\mathcal{H}}_i^{(\ell -1)}+\sum _{j\in \mathcal{N}\left(i\right)}{\mathcal{H}}_j^{(\ell -1)}\right),$$

(4)

where $\mathbf{w}$ is a learnable weight.

GIN has been proven to possess expressive power equivalent to the 2-Weisfeiler-Lehman test [58]. By utilizing GIN, we can effectively extract informative features from the superpixel graph, enabling accurate and efficient segmentation performance.

1.4. Memory Efficient Metric Learning

Although pixel-wise contrast can learn the global context to form a good segmentation embedding space [56], computing the contrastive loss requires using training image pixels, which leads to a significant amount of computation and memory cost. In this study, we propose superpixel-based methods that can significantly reduce the number of data samples from N to K, and we introduce a memory-efficient distance-based metric loss function.

The fundamental concept of metric learning is to bring similar samples closer together in the embedding space while pushing dissimilar samples further apart. However, pixel-wise contrast methods that involve setting a large number of anchor pixels and using tensor multiplication to compute positive similarity incur high memory costs. To address this, we define the superpixel metric loss using the mean square error (MSE) between the similarity and metric score of the embeddings, as follows:

$$\begin{array}{cc}\hfill & {\mathcal{L}}_{\mathrm{SM}}(\mathcal{A},\mathcal{R})=MSE(SIM(\mathcal{A},\mathcal{R}),Metric(\mathcal{A},\mathcal{R})),\hfill \end{array}$$

(5)

$$\begin{array}{cc}\hfill & SI{M}_{i,j}(\mathcal{A},\mathcal{R})=a·cos({\mathcal{A}}_i,{\mathcal{R}}_j)+b,SIM\in {\mathbb{R}}^{n_1\times n_2},\hfill \end{array}$$

(6)

$$\begin{array}{cc}\hfill & Metri{c}_{i,j}(\mathcal{A},\mathcal{R})=\frac{\mathcal{Y}\left({\mathcal{A}}_i\right)+\mathcal{Y}\left({\mathcal{R}}_j\right)}{2},Metric\in {\mathbb{R}}^{n_1\times n_2},\hfill \end{array}$$

(7)

where $n_1$ and $n_2$ represent the number of anchor samples $\mathcal{A}\in {\mathbb{R}}^{n_1\times d}$ and reference samples $\mathcal{R}\in {\mathbb{R}}^{n_2\times d}$, respectively. Here, d denotes the dimension of the embedding, and a and b are learnable parameters. Additionally, $\mathcal{Y}\left(x\right)$ represents the superpixel label of node x, which is defined by Eq. (1). It’s important to note that the anchor/reference samples are not restricted to being from the same image $\mathcal{A}$. The objective of Eq. (5) is to bring the embeddings of similar superpixel samples closer together and push dissimilar ones apart.

2. The Work-Flow and Architecture of SMGNN

We use the parameter-free methods to generate superpixels [40,59], on which we construct graph structure with two alternative strategies. Each superpixel is treated as a node and the mean RGB values and mean postion value consist five dimension node features. Suppose the mean scale of superpixels $S=\frac{H\times W}{K}$, we can define the adjacency between nodes according to their positional relation as

$${\mathcal{E}}_{i,j}=\left\{\begin{array}{cc}1,\hfill & |x_i-x_j|+|y_i-y_j|<\alpha \sqrt{S}\hfill \\ 0,\hfill & \mathrm{otherwise}\hfill \end{array}\right.,$$

where $(x_i,y_i)$ and $(x_j,y_j)$ are the positions of superpixel i and superpixel j respectively, and $\alpha \ge 2$ is a hyper-parameter to control the number of neighbouring nodes. The above definition will pose strong local connectivity to the graph data. For batched images, we can use parallel computing to accelerate the graph generation process and then multiple subgraphs combine as a large graph, where only connected nodes can perform message passing with the graph neural networks.

Regarding the model architecture, we utilize three layers of Graph Isomorphism Network to generate the embeddings of superpixels, upon which metric learning is performed. In addition to the transformed graph data derived from superpixels, we retain the original image data. We concatenate the features of superpixels and pixels and pass them through a lightweight convolutional neural network (CNN) to smooth out small pixel groups in the output of the GNNs. This process helps enhance the segmentation accuracy by incorporating non-degraded image information. The overall architecture of our proposed model, named SMGNN, is illustrated in Figure 3.

To tackle the clinical image segmentation, we employ the Dice loss [60], which is a structure-aware and widely used loss function for medical image segmentation. This loss function is designed to measure the similarity between predicted and ground truth segmentation masks. We also use the Dice coefficient to evaluate the performance of different models [60], which a widely used metric on image segmentation. To be more specific, given a set G, we define its characteristic/label function by ${\iota }_G\left(i\right)=\left\{\begin{array}{cc}\hfill 1,& i\in G\hfill \\ \hfill 0,& o.w.\hfill \end{array}\right.$. The Dice coefficient of two sets G and $\widehat{G}$ is defined as

$$Dice(G,\widehat{G})=\frac{2{\sum }_{i\in \mathsf{\Omega }}{\iota }_G\left(i\right)·{\iota }_{\widehat{G}}\left(i\right)}{{\sum }_{i\in \mathsf{\Omega }}\left({\iota }_G\left(i\right)+{\iota }_{\widehat{G}}\left(i\right)\right)},$$

(8)

where $\mathsf{\Omega }$ indicates the domain containing the two sets. The Dice metric is also directly used as loss function to train a supervised segmentation task. The Dice loss function is formulated as

$${\mathcal{L}}_{\mathrm{Dice}}(G,\widehat{G})=-\frac{2{\sum }_{i\in \mathsf{\Omega }}{\iota }_G\left(i\right)·{\iota }_{\widehat{G}}\left(i\right)}{{\sum }_{i\in \mathsf{\Omega }}\left({\iota }_G\left(i\right)+{\iota }_{\widehat{G}}\left(i\right)\right)},$$

(9)

The joint loss function for both the superpixel metric learning and segmentation tasks is defined as a combination of the Dice loss and the superpixel metric loss. This joint loss function enables us to optimize the model parameters simultaneously for both tasks, effectively leveraging the benefits of both superpixel-based metric learning and pixel-wise segmentation. The joint loss function is formulated as

$${\mathcal{L}}_{\mathrm{Joint}}={\mathcal{L}}_{\mathrm{Dice}}+\lambda {\mathcal{L}}_{\mathrm{SM}},$$

(10)

where $\lambda \ge 0$ is a hyperparameter to trade off the learning on image space and graph space.

3. Numerical Experiments

We employed a widely used White Blood Cells dataset to evaluate the effectiveness of our proposed recognizing system. We employed the widely-acclaimed Adam optimization technique [61] for model training during the backpropagation phase. The implementations are programmed with PyTorch-Geometric (ver 2.2.0) and PyTorch (version 1.12.1) and executed on an NVIDIA^® Tesla A100 GPU with $6,912$ CUDA cores and 80GB HBM2 installed on an HPC cluster.

3.1. Dataset Description

We verify the robustness of our methods on a small white blood cell image dataset collected from the CellaVision blog 1, which consists of one hundred 300×300 color images (30 neutrophils, 12 eosinophils, 3 basophils, 18 monocytes and 37 lymphocytes). The WBCs dataset leverages three-channel RGB images, which are processed via neural networks in an end-to-end training regimen. Each white blood cell image is manually-labeled, marking three primary regions: nuclei (represented in white), cytoplasm (depicted in grey), and the surrounding peripheral blood (captured in black). The white blood cell type is also A visual representation can be seen in Figure 5. The number of training/validation/testing data is 80%/10%/10% and the Dice loss function is applied. The dataset comprises five different cell types, staining effect, and illumination conditions which causes large variations in the sample distribution.

3.2. Evaluation of Superpixel Scale

To efficiently over-segment our input images, we adopted the SNIC superpixel generation methodology [59]. A distinctive feature of SNIC is its ability to visit each pixel just once—with the sole exception being those situated on superpixel boundaries. The computational traversal is characterized by the total number of pixels, N, augmented by a variable dictated by the desired superpixel count, K. Such a design renders SNIC more computationally nimble compared to alternatives like SLIC [40].

The superpixel quality, and its potential ramifications on segmentation, is an aspect we delve deeply into. By modulating the number of superpixels, we could ascertain its influence. For instance, the WBCs dataset results (illustrated by the red trajectory in Figure 4) signify that as the granularity of the superpixel method amplifies, there’s a corresponding upswing in the mIoU (mean Intersection over Union) score. Balancing optimal segmentation outcomes with computational practicality, we’ve pegged the mean scale of superpixels (S) at 16 for all ensuing experiments.

3.3. Comparison with Mainstream Deep Learning Segmentation Methods

A particularly intricate aspect of the WBCs dataset segmentation is the differentiation between the cytoplasm and the nuclei. Several images present nuclei that are sub-optimally stained, leading to a coloration reminiscent of the cytoplasm. This color overlap often ensnares traditional CNN-based segmentation models, resulting in predicted segmentations marred by interruptions or holes. Ideally, accurate segmentation demands that the representation of the nucleus be a cohesive, uninterrupted region. Our introduced SMGNN model takes an innovative approach by bolstering the connectivity of adjacent superpixels. Therefore, it is proficient in preserving the integrity of the nucleus region, where the capability is vividly showcased in Figure 5.

Figure 4. (a) The IoU reconstruction score VS number of superpixels. (b) Ground truth label image. (c)(d)(e)(f) Superpixel images and reconstructed label images (RLIs) of two different superpixel numbers. With an increase in the number of superpixels, the RLI tends to converge towards the ground truth label image. This trend indicates that the boundaries of the superpixels become more consistent with the boundaries of the cells, leading to improved quality of the superpixel segmentation.

Figure 4. (a) The IoU reconstruction score VS number of superpixels. (b) Ground truth label image. (c)(d)(e)(f) Superpixel images and reconstructed label images (RLIs) of two different superpixel numbers. With an increase in the number of superpixels, the RLI tends to converge towards the ground truth label image. This trend indicates that the boundaries of the superpixels become more consistent with the boundaries of the cells, leading to improved quality of the superpixel segmentation.

In Figure 6, we show the performance and number of learnable parameters of different deep learning models. Our SMGNN model can reach the SOTA performance while using far fewer parameters.

Figure 5. The segmentation results on the WBCs dataset. The ground truth (GT) annotation image contains a connected nuclei region without holes inside. Most CNN based methods tend to over-segment the nuclei region induced by the model bias. The SMGNN model can well preserve local connectivity and achieve comparable performance.

Figure 5. The segmentation results on the WBCs dataset. The ground truth (GT) annotation image contains a connected nuclei region without holes inside. Most CNN based methods tend to over-segment the nuclei region induced by the model bias. The SMGNN model can well preserve local connectivity and achieve comparable performance.

Figure 6. The comparison of model performance and network parameter size across different models on the WBCs dataset. The center of the circle indicates the Dice score of the model. The radius of the circle indicates the number of learnable parameters. The SMGNN model, utilizing approximately 7000 parameters, achieves comparable performance to models with millions of parameters.

Figure 6. The comparison of model performance and network parameter size across different models on the WBCs dataset. The center of the circle indicates the Dice score of the model. The radius of the circle indicates the number of learnable parameters. The SMGNN model, utilizing approximately 7000 parameters, achieves comparable performance to models with millions of parameters.

3.4. Ablation Study

To optimize node embeddings within our methodology, we selected the GIN model for the GNN module due to its superior discriminative capacity. Comparative tests with popular GNN models like GCN and GAT revealed that a configuration using three layers of GIN demonstrated an enhanced performance, significantly improving node classification accuracy.

Beyond the conventional setup detailed in Figure 3, which incorporates pixel-level embedding, we ventured into an approach that solely leverages a GNN, transitioning from image-level segmentation components to a strictly node-based classification method, as illustrated in Figure 7. The training for this node classification is driven by a cross-entropy loss function, delineated as follows:

$${\mathcal{L}}_{\mathrm{CE}}(\mathcal{S},\widehat{\mathcal{S}})=-\frac{1}{K}\sum _{i=1}^K\sum _{c=1}^C{\mathcal{S}}_{ic}log\left({\widehat{\mathcal{S}}}_{ic}\right),$$

(11)

where we use the majority voting rule to define the supepixel label as

$$\mathcal{S}=\lfloor \mathcal{Y}+0.5\rfloor ,$$

(12)

guiding the supervised learning in graph space. In Eq. (12), $\lfloor ·\rfloor$ is the round down function, and $\widehat{\mathcal{S}}$ is the predicted probability of superpixel, and C is the number of semantic classes. Though such a rule may group mistaken pixels whose pixel label is not the majority, the pure GNN methods may achieve good performance when the scale of the superpixel is small. Our ablation studies, as depicted in Figure 8, offer keen insights into the performance nuances of pure GNN-based segmentation models. These models manifest commendable segmentation outcomes when oriented to small-scale superpixel environments. However, as we scale the superpixels, the model’s performance is inversely impacted by its heightened sensitivity to superpixel quality, leading to notable performance drops. Introducing convolutional filters via CNN feature embedding for image-level segmentation does augment the model with additional parameters. Nevertheless, the significance of these filters is evident in the stability they confer upon the model, especially when navigating varying superpixel scales.

This investigation underpins a critical takeaway: the scale of superpixels and a model’s sensitivity to their quality must be harmoniously calibrated. Relying exclusively on GNN-driven segmentation models may prove suboptimal when maneuvering larger superpixel frameworks.

3.5. Effectiveness of Metric Learning on Embedding Space

To understand the impact of metric learning on the embedding space, we visually represent the spatial relationships of superpixels. We assign different colors to these superpixels based on their labels, as determined by Eq. (12). Figure 9, created using UMAP [62], demonstrates the distances among embeddings derived from the GNN. A notable observation from this visualization is the pronounced separation between superpixels with distinct labels—a testament to the efficacy of incorporating metric learning. Furthermore, there’s a heightened cosine similarity between samples that are alike, while distinct samples exhibit reduced similarity. This distinction underscores the model’s ability to effectively differentiate and group superpixels in the embedding space.

3.6. White Blood Cells Type Classification Results

Though the segmentation of cell salient regions, such as the nucleus and cytoplasm, is fundamental and challenging, there are various off-the-shelf methods available for subsequent cell type classification [63,64,65]. Segmentation may provide a direct means to obtain distinguishing characteristics for cell type classification, as cell morphology is closely related to type.

In this research, we employ a light-eight ResNet neural network [66,67] and train a classifier based on the output of segmentation network, as shown in Figure 10. We sample about $1/3$ training images from each classes and leave the remaining for testing. The classification result is shown on Table 1, and our classification method can achieve about $96.72\%$ overall accuracy. In addition to our proposed segmentation-based cell type recognition system, we also implemented a baseline method to predict cell types without utilizing segmented cell regions. The corresponding results are presented in Table 2, and such method without extraction on slient cell regions can barely achieve $72.13\%$ overall accuracy. There are also traditional methods employ hand-crafted features extracted from segmented regions, combined with machine learning classifiers like Support Vector Machine (SVM) [63], and such method can get overall accuracy ranging from $89.69\%$ to $96\%$. Our proposed deep learning-based automatic recognition system demonstrates high efficiency, and its accuracy can be further improved with an increase in the number of available training samples.

4. Conclusion

In this research paper, we propose a deep learning based automatic recognizing system for the challenging white blood cell image recognizing task. In the first part, our proposed SMGNN segmentation model, which combines superpixel methods and a lightweight Graph Isomorphism Network to significantly reduce memory usage while preserving segmentation capabilities. We innovatively propose superpixel metric learning to capture cross-image global context information, making it highly suitable for medical images with limited training samples. Comparing our model to other mainstream deep learning models, we achieve comparable segmentation performance with a remarkable reduction of at most 10000 times fewer parameters. Through extended numerical experiments, we further investigate the effectiveness of metric learning and the quality of superpixels In the second part, the segmentation-based cell type classification processes exhibit satisfactory results, indicating that the overall automatic recognition algorithms are accurate and efficient for execution in hematological laboratories. We have made our code publicly available at https://github.com/jyh6681/SPXL-GNN, and we encourage its widespread implementation in portable devices of hematologists and remote rural areas.