Code accompanying the paper "Fighting Class Imbalance with Contrastive Learning" (MICCAI2021)
https://lmb.informatik.uni-freiburg.de/Publications/2021/MMB21/
Yassine Marrakchi, Osama Makansi, Thomas Brox.
The algorithm is based on supervised contrastive loss. It is able to efficiently learn good features for imbalanced image classification. The code supports the classification for automated diagnosis of pigmented skin lesions using the HAM10000 dataset. The used dataset can be downloaded from: https://www.kaggle.com/kmader/skin-cancer-mnist-ham10000
Pretraining stage:
python main_supcon.py --batch_size 192 \
--data_folder /path/to/dataset \
--balance \
--cosine
Removing the flag --balance
would lead to learning from the original unbalanced data.
Linear evaluation stage:
python main_linear.py --batch_size 512 \
--data_folder /path/to/dataset \
--ckpt /path/to/model.pth
Results on HAM-10000 (also used in ISIC-2018 lesion diagnosis challenge):
@inproceedings{DBLP:conf/miccai/MarrakchiMB21,
author = {Yassine Marrakchi and
Osama Makansi and
Thomas Brox},
title = {Fighting Class Imbalance with Contrastive Learning},
booktitle = {Medical Image Computing and Computer Assisted Intervention - MICCAI},
year = {2021},
doi = {10.1007/978-3-030-87199-4\_44},
}