A Classification Method for Optical Coherence Tomography Images Based on a Structure-oriented Adaptive Neural Network Architecture

Abstract

The method of optical coherence tomography image classification for automated diagnosis of diabetic retinopathy and diabetic macular edema is proposed in the article. An innovative adaptive multi-task deep neural network is created. It simultaneously solves the problems of pathology classification and structural feature reconstruction. The neural network uses the pre-trained EfficientNetB7 model as an encoder for efficient extraction of high-level features. The structural feature learning branch (decoder) is responsible for restoring spatial information. It increases the resolution of feature maps to the original size of 224x224 pixels with a gradual decrease in the number of filters and the use of Batch Normalization to stabilize learning. The classification branch combines semantic and structural features. It uses the channel attention mechanism for dynamic weighting of informative channels. Dropout and Batch Normalization layers are used to prevent overtraining in the classification branch. The model is optimized using a multi-task loss function. It consists of a modified loss function for classification (with class weights to balance data imbalance) and a root-mean-square error for structural loss. Training is performed using the Adam optimizer and the EarlyStopping, ModelCheckpoint, and ReduceLROnPlateau callbacks. The experiment was conducted on the OCT Image Classification dataset. Data augmentation (horizontal reflections) was performed to increase the number of images. High accuracy rates and cost functions were obtained as a result of training. The multi-task method enables the encoder to learn details and boundaries of the retina through Canny edge reconstruction. It contributes to improved classification and provides a powerful internal regularization mechanism, increasing the generalization ability of the model.