Deep Colorization modules for Medical Image Analysis

This project aims at bridging the gap between medical image analysis by introducing a light colorization module. Three different modules are proposed and implemented (DECONV, PixelShuffle and ColorU)

The modules are trained jointly with a backbone pre-trained on ImageNet. A multi-stage transfer learning pipeline is summarized here. First, the colorization module is trained from scratch together with the classifier, while the pre-trained CNN backbone is kept frozen, to learn the mapping which maximizes classification accuracy. Then, the entire network is fine-tuned to learn useful features for the target task, while simultaneously adjusting the colorization mapping. The figure below shows the output of each colorization module when only the colorization module is trained, and after the entire network is fine-tuned.

Dependencies

Linux
Python 3.7
PyTorch 1.4.0

Download

Trained models with DenseNet121 and ResNet18 backbones are available here A detailed list of models is available here

All models were trained on CheXpert to predict the presence/absence of 5 labels:

Atelectasis
Cardiomegaly
Consolidation
Edema
Pleural Effusion

Image normalization

If you wish to use the above models, please bear in mind that images were normalized with statistics calculated on the CheXPert dataset:

mean: [0.5028, 0.5028, 0.5028]
std: [0.2902, 0.2902, 0.2902]

Citation

If you use the models in your research, please cite our paper:

@article{morra2020bridging,
  title="Bridging the gap between Natural and Medical Images through Deep Colorization",
  author="Morra, Lia and Piano, Luca and Lamberti, Fabrizio and Tommasi, Tatiana",
  year="2020"
  }