W-net: Simultaneous segmentation of multi-anatomical retinal structures using a multi-task deep neural network

TL;DR

This paper introduces a multi-task deep neural network called W-net for simultaneous segmentation of optic disc and exudates in retinal images, improving accuracy and robustness through specialized loss functions and cross-dataset validation.

Contribution

The study presents a novel multi-task learning framework with class-balanced and multi-task weighted losses for joint segmentation of multiple retinal structures, outperforming separate models.

Findings

Achieved high F1-scores of 94.76% and 95.73% for OD segmentation.

Achieved F1-scores of 92.80% and 94.14% for exudates segmentation.

Demonstrated improved generalization across multiple datasets.

Abstract

Segmentation of multiple anatomical structures is of great importance in medical image analysis. In this study, we proposed a $\mathcal{W}$-net to simultaneously segment both the optic disc (OD) and the exudates in retinal images based on the multi-task learning (MTL) scheme. We introduced a class-balanced loss and a multi-task weighted loss to alleviate the imbalanced problem and to improve the robustness and generalization property of the $\mathcal{W}$-net. We demonstrated the effectiveness of our approach by applying five-fold cross-validation experiments on two public datasets e\_ophtha\_EX and DiaRetDb1. We achieved F1-score of 94.76\% and 95.73\% for OD segmentation, and 92.80\% and 94.14\% for exudates segmentation. To further prove the generalization property of the proposed method, we applied the trained model on the DRIONS-DB dataset for OD segmentation and on the MESSIDOR dataset for exudate segmentation. Our results demonstrated that by choosing the optimal weights of each task, the MTL based $\mathcal{W}$-net outperformed separate models trained individually on each task. Code and pre-trained models will be available at: \url{https://github.com/FundusResearch/MTL_for_OD_and_exudates.git}.