ProCo: Prototype-aware Contrastive Learning for Long-tailed Medical Image Classification
It addresses the problem of imbalanced medical image classification for healthcare applications, presenting a novel single-stage pipeline as an incremental improvement over existing two-stage methods.
The paper tackles long-tailed medical image classification by proposing ProCo, a prototype-aware contrastive learning framework that generates representative contrastive pairs and recalibrates prototypes to address data imbalance, achieving state-of-the-art performance with large margins on two highly-imbalanced datasets.
Medical image classification has been widely adopted in medical image analysis. However, due to the difficulty of collecting and labeling data in the medical area, medical image datasets are usually highly-imbalanced. To address this problem, previous works utilized class samples as prior for re-weighting or re-sampling but the feature representation is usually still not discriminative enough. In this paper, we adopt the contrastive learning to tackle the long-tailed medical imbalance problem. Specifically, we first propose the category prototype and adversarial proto-instance to generate representative contrastive pairs. Then, the prototype recalibration strategy is proposed to address the highly imbalanced data distribution. Finally, a unified proto-loss is designed to train our framework. The overall framework, namely as Prototype-aware Contrastive learning (ProCo), is unified as a single-stage pipeline in an end-to-end manner to alleviate the imbalanced problem in medical image classification, which is also a distinct progress than existing works as they follow the traditional two-stage pipeline. Extensive experiments on two highly-imbalanced medical image classification datasets demonstrate that our method outperforms the existing state-of-the-art methods by a large margin.