How many labels do I need? Self-supervised learning strategies for multiple blood parasites classification in microscopy images

2.1. Dataset

An extensive multi-center clinical dataset of blood samples from patients suspected of parasitic infection (filariae, malaria, chagas disease, and leishmaniasis) was collected and digitized with a smartphone attached to a conventional microscope using a 3D-printed adapter (Spotlab, Spain) [13]. Blood samples from 332 subjects were digitized, yielding 11,956 microscope field-of-view images (FoV). To increase the variability of our dataset, we included a public dataset of Chagas with 675 images [7], thus forming a large database with images from 4 different sites. Our database includes images with different magnifications: 10× (∼0.81 µm/pixel), 40× (∼0.20 µm/pixel), 100× (∼0.081 µm/pixel), each FoV measures ∼2700 x 2700 pixels.

To tailor the dataset for our classification objectives and maintain granularity, we divided each FoV into a structured 3×3 grid layout, resulting in a total of 105,075 cropped patches that were resized to 300x300 pixels. FoV images were labeled by expert microscopists. From those images, we curated a subset of 15,268 labeled patches that contained parasites across 11 unique classes, as detailed in Table 1.

These 15,268 images were used for training and assessing model classification performance. The annotated images were further divided into training and validation sets using a 5 fold cross-validation split at patient level and stratified by species distribution. On the other hand, the remaining 89,600 unlabeled patches were used for training the self-supervised phase of the proposed approach. To maintain the integrity of the training process, the labeled dataset for training the classifier and the unlabeled dataset were segregated so that no image patches are used in both the self-supervised phase and the subsequent classification.

2.2. Overview of the proposed method

Utilizing large-scale pretrained models has been proven to offer diverse advantages, such as enhancing downstream task performance and minimizing the required labeled data for subsequent training tasks [14], [15]. Figure 1 illustrates the proposed pipeline. To summarize, we first initialize a backbone encoder network with weights from ImageNet, followed by transferring these weights to train the SSL network. During this phase, the network undergoes training using a large, unlabeled dataset of microscopy blood images, allowing it to autonomously acquire visual representations of blood parasites without the need for manual annotations. Subsequently, transfer learning is once again applied, culminating in the network’s final training for the downstream task using the labeled dataset. The primary objective of this final network is to classify each microscopy image, discerning the presence of the previously defined 11 possible blood parasites. This streamlined pipeline optimizes the advantages of pretrained models and self-supervised learning, ensuring robust performance while minimizing the demand for extensive labeled data.

• Download figure
• Open in new tab

Fig. 1.

Proposed methodology, comprising: (1) supervised representation learning on a large-scale dataset of labeled natural images such as ImageNet; (2) self-supervised contrastive representation learning on a large unlabelled microscopy dataset; and (3) supervised fine-tuning on manually labeled dataset for downstream task.

2.3. Image representation with self-supervised learning

We used SSL for improving blood microscopy image representation by training a feature extractor (backbone, ResNet50) in order to increase data generalization, by leveraging unlabeled datasets. In this work, we implemented a method based on SimSiam [16], a non-contrastive SSL algorithm that, in contrast to SimCLR [17] or MoCO [18], does not need negative samples nor large batches and tries to optimize the similarity among images using only positive pairs. In order to learn meaningful representations in a self-supervised manner, we used a siamese network which tries to maximize the similarity between two different views (transformations) of the same original image. For this purpose, given an image x, two different views (x₁, x₂) are generated with random augmentations T and T’. Such transformations hold the potential to influence the efficacy of a pre-trained model increasing its generalization capability. The transformations selected for our investigation include random cropping, color jittering, and random flipping, all of which align with the augmentation strategies utilized within the SimCLR framework [17] and are commonly adopted as baseline augmentations in a wide array of SSL algorithms. Hyperparameters of these transformations (e.g. crop scale or brightness) were selected with careful consideration of the clinical context of microscopy images, aiming to prevent the generation of non-realistic results as well as to ensure color and space invariance. The generated views are processed by an encoder, formed by a convolutional backbone and a 3-layer multi layer perceptron (MLP) projector to generate image embeddings z₁and z₂. Then, each embedding, z, is processed by a 2-layer bottleneck MLP predictor producing embeddings p₁ and p₂. Loss function was defined based on the negative cosine similarity between the vectors z₁ and p₂:

A key factor of this algorithm is the introduction of stop gradients to avoid representation collapse as seen in Figure 1 By updating only the gradient of the network output taking into account the output of the predictor, it avoids the production of constant value for every image leading to no improvement in the representation learning [19]. This process is done once per each pair of views(x₁, x₂) leading to the following symmetrized loss function:

2.4. Incremental training for parasite classification

Once the backbone architecture is trained via self-supervised learning, weights are transferred and fine-tuned for the downstream classification task. At this point, we tested two different approaches. The first approach was based on only allowing the training of the weights of the last convolutional layers and those from the classification dense layers, and leaving the weights of the first convolutional layers frozen (linear probe). In the second approach, we allow fine-tuning all weights of the architecture.

To assess the impact of self-supervised learning (SSL) in a constrained labeled dataset scenario, we conducted an incremental training experiment, augmenting the training size in steps of 5%, 10%, 15%, 25%, 50%, 75% and 100%.

Employing 5-fold cross-validation, we systematically evaluated parasite classification performance across 11 classes, comparing the SSL pretrained backbone to a baseline ResNet50 classifier with ImageNet weights.

All classification networks were implemented to optimize the cross-entropy loss. During training, various data augmentation techniques were applied such as random flipping, brightness and saturation transformations. To handle the imbalanced distribution across the 11 classes of our labeled dataset and to reduce potential overfitting to the majority classes, we weighted the loss function according to class distributions.

3.1. Self-supervised pre-training

ResNet50 backbone was initialized with ImageNet weights in order to speed up the pre-training. The previously described loss function for training the SSL network was optimized with SGD method with 1e-4 weight decay and 0.9 momentum, with a cosine decayed learning rate of 0.01× (batch size / 256). We used a batch size of 32 and pre-trained the model over a span of 25 epochs.

Figure 2. reveals the model’s capability to extract relevant features illustrating a high degree of intra-class similarity and a significant inter-class separation in the feature latent space. This shows that the learned representations are meaningful to distinguish the multiple parasite species.

• Download figure
• Open in new tab

Fig. 2.

tSNE visualization of the features extracted by SimSiam.

3.2. Incremental training and classification results

We trained with the Adam optimizer with 1e-4 weight decay, with an initial learning rate of 0.01 with Cosine Decay. All models were evaluated on the same validation set, and the performance was assessed using the macro (unweighted) F1 score. Figure 3 shows that leveraging the foundational SSL weights enhances performance compared to training from scratch, particularly in scenarios with limited labeled data. Table 2 provides a breakdown of model performance, including SSL performance when fine-tuning all weights rather than just the last layers.

• Download figure
• Open in new tab

Fig. 3.

Comparative performance (macro F1-score) analysis illustrating the impact of employing the SSL methodology versus not using it, across an incremental training experiment.

Table 2 illustrates that freezing initial weights (linear probe) when little labeled data is available significantly boosts performance. Full-model fine-tuning with scarce data risks losing vital pretrained features and may lead to overfitting. In contrast, restricting fine-tuning to the latter layers markedly improves performance, utilizing stable features from self-supervised pre-training for task adaptation. Detailed performance of our proposed method for each class under study is shown in Figure 4, where we show that our model with just ∼100 labels can achieve a mean F1 score of 0.798 and a deviation of 0.074 across the 11 parasite species within our dataset.

• Download figure
• Open in new tab

Fig. 4.

F1-Score comparison for each class using the SSL model.

Figure 5. provides a detailed depiction of the outcomes presented in the previous figure, specifically highlighting our model performance across individual classes when the number of training labels is close to 100.

• Download figure
• Open in new tab

Fig. 5.

F1-Score when training labels is ∼100 for each class

In order to assess our SSL model’s generalization capability, we tested it on an independent malaria blood sample classification task using a public database of 2,025 images (1,751 for training, 274 for testing) across 7 classes (various malaria species and maturation stages) [20]. Transferring weights from our SSL model improved classification (macro F1-score), from 38.90% (using a ResNet50 initialized with ImageNet weights) to 83.18%.