DAug: Diffusion-based Channel Augmentation for Radiology Image Retrieval and Classification

TL;DR

DAug leverages diffusion-based generative models to augment radiology images with disease region heatmaps, enhancing AI performance in medical image retrieval and classification without altering existing architectures.

Contribution

The paper introduces a novel diffusion-based feature augmentation method and a hybrid contrastive learning approach, significantly improving medical image understanding tasks.

Findings

Achieves state-of-the-art results on medical image retrieval.

Enhances classification accuracy with heatmap augmentation.

Improves robustness of AI models in radiology analysis.

Abstract

Medical image understanding requires meticulous examination of fine visual details, with particular regions requiring additional attention. While radiologists build such expertise over years of experience, it is challenging for AI models to learn where to look with limited amounts of training data. This limitation results in unsatisfying robustness in medical image understanding. To address this issue, we propose Diffusion-based Feature Augmentation (DAug), a portable method that improves a perception model's performance with a generative model's output. Specifically, we extend a radiology image to multiple channels, with the additional channels being the heatmaps of regions where diseases tend to develop. A diffusion-based image-to-image translation model was used to generate such heatmaps conditioned on selected disease classes. Our method is motivated by the fact that generative models learn the distribution of normal and abnormal images, and such knowledge is complementary to image understanding tasks. In addition, we propose the Image-Text-Class Hybrid Contrastive learning to utilize both text and class labels. With two novel approaches combined, our method surpasses baseline models without changing the model architecture, and achieves state-of-the-art performance on both medical image retrieval and classification tasks.