Progressive Attention Guidance for Whole Slide Vulvovaginal Candidiasis Screening

We use a pre-trained detection model as prior to initialize the classification model. In experimental exploration, we find that, if we train the detection network directly, the bounding-box annotation indicates the location of candida and can rapidly establish a rough understanding of the morphology of candida. However, the positioning task coming with detection lacks enough granularity, resulting in relatively low precision to discern cell edges or folding from candida. Meanwhile, directly training a classification model is usually easier to converge. However, in such a task, as candida occupies only a few pixels in an image, it is difficult for the classifier to focus on the target. That, the attention of the classifier may spread across the entire image, leading to overfitted training quickly.

Therefore, we argue that the detection and classification tasks are complementary to solve our problem. Particularly, we pre-train a detector and inherit its advantages in the classifier. We use Retinanet[22], which is composed of a backbone attached with FPN (Feature Pyramid Network, FPN, [23]) and a detection head, as shown in (Fig. 3). We chose the same encoder architecture (Resnet[21]) for the detection and classification networks. To train the encoder with the detection task (Fig. 3), we use bounding-box annotations to supervise Retinanet. We then initialize the classification network by directly loading the encoder parameters and freezing the first few layers during the training of the classification network. Note that pre-training not only discards the complex positioning task but also makes it easier for the classification network to converge especially in the early stage of training.

We design a novel skip self-attention (SSA) module to fuse discriminative features of candida from different scales. At a fine-grained level, the hyphae and spores of candida are usually the basis for judging. Yet we need to distinguish them from easily distorting factors such as contaminants in WSIs and edges of nearby cells. At a coarse-grained level, there is the phenomenon that a candida usually links multiple host cells and yields a string of them. Thus it is necessary to combine long-range visual cues that span several cells to derive the decision related to candida.

CNN-based methods have achieved excellent performance in computer-aided diagnosis including cervical cancer [19]. However, the unique shape and appearance of the candidate incur troubles for CNN-based classifiers, whose spatial field of view can be relatively limited. In recent years, vision transformer (ViT) has been widely used in visual tasks for its global attention mechanism [12], sensitivity to shape information in images [11], and robustness to occlusion [10]. Nevertheless, such a transformer can be hard to train for our task, due to the large image size, huge network parameters, and huge demand for training data. Therefore, to adapt to the shape and appearance of candida, we propose the SSA module and apply it to ViT for efficient learning.

Specifically, we use the pre-trained CNN-based encoder to extract features for each cropped image. The feature maps extracted after the first layer is considered low-level, which contains fine-grained texture information. On the contrary, the feature maps extracted from the last layer are high-level, which represents semantics regarding candida. To combine the low- and high-level features, we regard the low-level features as queries (Q), and the high-level features as keys (K) and values (V). For each patch in the ViT scheme, the transformer decoder computes the attention between low- and high-level features and combines them. The class token ’CLS’ is used for the final classification. The combined feature maps can offer more representative information so that the classifier focuses more on different scales to long-range candida. Meanwhile, the extra SSA structure is simple, which causes a low computation burden.

As mentioned in Section 1, the style gap is another problem, which makes overfitting more severe. In this part, we adopt the strategy of contrastive learning to alleviate such problems and further optimize the attention of the network. Our approach has two key goals: (1) to ensure that the features from the original image remain consistent after undergoing various image augmentations, and (2) to construct an image without the region of candida, resulting in highly dissimilar features compared to the original.

Inspired by a weakly supervised learning segmentation method[8], we construct a contrastive learning method, which will be described in detail in the following sections, as shown in Fig. 3.

To achieve this, we use augmentation and the attention map generated during the training process to construct three types of images and apply contrastive learning to the features extracted from them. For a given image $I$, we use image augmentation to generate $I_{aug}$ and use the encoder attached with SSA to extract feature, $F^{c}_{aug}$. The attention map $A$ is transformed from $F^{c}_{aug}$ by an attention extractor. The attention extractor uses FC (the same params as that of the classifier) to reduce the channels of features (except the class token) to 2 (candida and others), then reshape the features representing candida to a feature map, and applies bilinear interpolation to upsample it to the same size of $I$. Eq. 1 normalizes $A$ to the interval [0,1], obtaining $M$ to represent the likelihood of candida distribution. We get the masked image $I_{masked}$ by subtracting $M$ from $I$.

where $\sigma$ and $s$ are used to adjust the range of values, set to 0.5 and 10.

As shown in Fig. 3, the features $F$, $F_{aug}$ and $F_{masked}$ from the three types of images $I$, $I_{aug}$ and $I_{masked}$ by the shared classifier. In our task, we hope that the style gap does not affect the feature extraction of the image, so the distance between $F_{aug}$ and $F$ should be attracted. At the same time, we hope that $F_{masked}$ should not contain the characteristics of candida, which is repelled from $F_{aug}$. To achieve our goal, we introduce triplet loss[24] for contrastive learning as shown in the first part of Eq. 2.

In addition, we use two constraints, leading to more stable and robust training. If our network has effective attention, the masked image should not contain any candida, so the score of the Candida category after the mask $S(I_{masked})$ should be minimized. We use the attention mining loss to handle this, as shown in the second part of Eq. 2. Additionally, we need the attention to cover only the partial area around candida, without false positive regions. Otherwise, attention maps that cover the whole image can also result in low $L_{tri}$. We take the average grayscale of attention map $\bar{M}$ as a restriction, as shown in the last part of Eq. 2. $L_{tri}$,$L_{am}$, and $L_{focus}$ are combined as $L_{cl}$ to constrain each other and take full advantage of contrastive learning, as shown in Eq. 2.

Finally, we use the cross-entropy loss $L_{ce}$ to calculate the classification loss with $labels$. The total loss during training can be expressed as shown in Eq. 3. $\alpha$ is a hyper-parameter, set to 0.1.

With the strategies above, we have built the classifier for all cropped images in a WSI. Then we can finish the pipeline of WSI-level classification, which is shown as part of Fig. 2. Specifically, for each cropped image, we conduct image-level classification to find out whether it suffers from suspicious candida infection. The image-level classification also produces a score, as well as the feature representation of the image under consideration. Then, we reorganize features from all images by ranking their scores and preserving that with top-k scores. We complete the aggregation of the top-k features by the transformer and make the WSI-level decision by an FC layer in the final.