Attention-Based Modules for Diabetic Retinopathy Diagnosis

Introduction

Diabetic Retinopathy (DR) is one of the leading causes of blindness, and its early detection is critical to managing the progression of diabetes and avoiding severe complications. Computational approaches, especially based on fundus imaging, have shown great promise in diagnosing DR with high precision and speed.

In this project, we focus on improving DR diagnosis using attention-based modules in a semantic segmentation architecture, primarily a U-Net model. We integrate several attention mechanisms—Attention Gait (AG), Multihead Attention (MHA), Convolutional Block Attention (CBAM), and Squeeze Excitation (SE)—to enhance lesion detection (such as hemorrhages and microaneurysms) and improve classification performance. Our experiments are conducted on multiple datasets, including IDRID, FGADR, EyePacs, and APTOS. You can find the code here

Dataset Overview

The datasets used in this project provide labeled fundus images, indicating varying degrees of Diabetic Retinopathy. These datasets include:

• IDRID: Indian Diabetic Retinopathy Image Dataset, annotated with various lesion types such as hemorrhages, exudates, and microaneurysms.
• FGADR: Fundus Grading and Annotation for Diabetic Retinopathy, designed for DR lesion detection.
• EyePacs: A widely used dataset for binary classification of DR presence or absence.
• APTOS: Dataset from the Asia Pacific Tele-Ophthalmology Society challenge, consisting of DR-graded fundus images.

These datasets provide a wide range of images necessary for training and validating the models for both lesion detection and classification.

1. Semantic Segmentation and Attention-Based Models

In this project, we enhanced a U-Net-based semantic segmentation model by integrating different attention modules. These modules aim to focus on critical regions of the fundus images where DR-related lesions are present. Below is a brief overview of each module used:

Attention Modules:

• Attention Gait (AG): A module designed to prioritize spatial locations in the image by focusing attention on critical lesion areas.
• Multihead Attention (MHA): Derived from Transformer models, MHA allows the model to focus on different lesion-related features simultaneously.
• Convolutional Block Attention (CBAM): This module combines spatial and channel-wise attention, focusing on both where and what features to emphasize.
• Squeeze Excitation (SE): A mechanism that recalibrates channel-wise features by focusing on their importance for lesion detection.

EyePACS Dataset Preprocessing Summary

The preprocessing of the EyePACS dataset involved multiple stages to enhance image quality and diversify the data through augmentation. The key steps are summarized below:

1. Random Image Display

A function was developed to visualize random images from the dataset. This allowed for inspection of data loading and visualization before applying augmentations.

        def plot_random_image(folder_path):
            image_list = [filename for filename in os.listdir(folder_path) if filename.endswith('.jpeg')]
            random_image = random.choice(image_list)
            image_path = os.path.join(folder_path, random_image)
            image = cv2.imread(image_path)
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            plt.imshow(image_rgb)
            plt.axis('off')
            plt.show()
        

2. Image Augmentation Techniques

To increase the diversity of the dataset, several augmentation techniques were applied:

• Gaussian Blur
• Gamma Contrast Adjustment
• Image Sharpening
• Horizontal Flipping
• Rotation (-25° to 25°)

Example of applying Gaussian Blur:

        aug = iaa.GaussianBlur(sigma=(0.0, 1.0))
        image_aug = aug(image=image)
        ia.imshow(image_aug)
        

3. CLAHE (Contrast Limited Adaptive Histogram Equalization)

CLAHE was used to improve image contrast, especially for underexposed images.

        def aplicar_filtro_clahe(imagen):
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            imagen_clahe = clahe.apply(imagen)
            return imagen_clahe
        

4. Data Augmentation for Bright Images

A check was implemented to avoid augmenting images with low brightness. If the image brightness exceeded a threshold, augmentations like color multiplication, blur, and sharpening were applied.

        brightness_threshold = 84
        if brightness > brightness_threshold:
            augmented_image = augmenter.augment_image(image)
        

5. SNR (Signal-to-Noise Ratio) Analysis

The Signal-to-Noise Ratio (SNR) was calculated for each image. Based on the SNR value, images were either included for augmentation or excluded if they had too much noise.

        snr = 20 * np.log10(mean / std_dev)
        if snr >= threshold:
            # Proceed with augmentation
        

6. Blurring Index

The blurring index was computed using the Laplacian variance to determine the focus quality of images.

        estimator_operator = cv2.Laplacian(image, cv2.CV_64F).var()
        blurring_index = estimator_operator / std_deviation
        

IDRiD Dataset Preprocessing Summary

The preprocessing for the IDRiD dataset included a U-Net model with the CBAM (Convolutional Block Attention Module) for segmentation tasks. Below is a summary of the key steps:

1. U-Net Model Architecture

The U-Net architecture consists of an encoder-decoder structure:

• **Encoder:** Series of convolutional layers with max-pooling for down-sampling.
• **Bottleneck:** Middle section with higher number of filters for feature extraction.
• **Decoder:** Up-sampling and concatenation of features from the encoder with skip connections.
• **Output:** Final layer with a sigmoid activation function for segmentation.

        def unet_model(input_shape, num_classes):
            inputs = Input(input_shape)
            conv1 = Conv2D(64, 3, activation='relu', padding='same')(inputs)
            pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
            ...
            outputs = Conv2D(num_classes, 1, activation='sigmoid')(conv9)
            model = Model(inputs=inputs, outputs=outputs)
            return model
        

2. CBAM (Convolutional Block Attention Module)

CBAM was integrated into the U-Net to improve feature selection by applying attention mechanisms both spatially and across channels.

• **Spatial Attention:** Focuses on important areas within the image.
• **Channel Attention:** Highlights important feature maps (channels).

        def cbam_block(inputs, ratio=8):
            channels = inputs.shape[-1]
            spatial_max = tf.reduce_max(inputs, axis=[1, 2], keepdims=True)
            ...
            return layers.Add()([spatial_attention, channel_attention])
        

3. Data Loading and Preprocessing

Images and corresponding masks were loaded and resized. The images were normalized, and a single channel was added to handle grayscale images.

        def load_images_and_masks(image_filenames, mask_dir, input_shape):
            images = []
            masks = []
            for image_filename in image_filenames:
                image_path = os.path.join(image_dir, image_filename)
                image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
                ...
            return np.array(images), masks
        

4. Model Compilation and Training

The U-Net model with CBAM was compiled using the Adam optimizer and sparse categorical cross-entropy loss. The model was trained using an 80/20 split for training and validation data.

        model.compile(optimizer=Adam(learning_rate=1e-4),
                      loss='sparse_categorical_crossentropy',
                      metrics=[mean_iou, mean_iou_thresholded()])
        

5. Custom Metrics: Mean IoU

Custom metrics were implemented to calculate the Intersection over Union (IoU), a crucial metric for segmentation tasks.

        def mean_iou(y_true, y_pred):
            y_pred = K.argmax(y_pred, axis=-1)
            ...
            return K.mean(iou)