Weakly supervised lesion localization and attribution for OCT images with a guided counterfactual explainer model

Abstract

Lesion localization plays an important role in computer aided diagnosis. Due to the lacking of lesion annotations, weakly supervised methods using only image-level annotations are demanded for a wide variety of applications, especially optical coherence tomography diagnosis. Most weakly-supervised methods rely on attribution analysis. However, current methods suffer from imprecise attributions and lead to poor localization quality. To address this problem, we introduce a lesion localization method based on a new explainable AI approach, termed Optical Coherence Tomography Class Association Embedding (OCT-CAE), leverages image-level annotations and a cycle generative adversarial network with subspace decomposition to fuse global knowledge and enable counterfactual generation. Each sample is encoded into a pair of subspaces where a low-dimensional common subspace is created to embed manifold of classification-related information and an individual subspace to embed individual-specific information. With the trained OCT-CAE, the codes in the two subspaces can be freely recombined to generate realistic images. These generated images retain the class-related features defined by the common subspace while preserving individual-specific characteristics from the individual subspace. Lesion localization is achieved by altering the common code to induce a class-flip in the generated images. By comparing the modified and original images, we can identify lesion regions without requiring regional annotations. Extensive experiments on two publicly available datasets demonstrate that OCT-CAE effectively disentangles latent information space in images, achieving state-of-the-art performance. Our code is available at https://github.com/lcmmai/OCT-CAE.