Task-Aware Transformer For Partially Supervised Retinal Fundus Image Segmentation

Hailong Zeng, Jianfeng Liu, Yixiong Liang

Abstract: The segmentation of retinal fundus images plays a critical role in ophthalmic disease screening and diagnosis. Due to the considerable human labor and expertise required for data annotation, fundus image datasets are often only partially labeled. Existing approaches typically employ separate networks for each specific task, resulting in resource-intensive processes and limited generalization capability. To address this, we propose a novel task-aware transformer (TAFormer) that learns to segment lesions and anatomic structures on multiple partially labeled datastes. Our method is built upon the query-based learning approaches. Specifically, we propose a task-specific query grouping strategy, which enables each group of queries to learn structural information related to the corresponding task. Queries from different groups can then capture relationships between the current task and other tasks through self-attention mechanisms, thereby enhancing their own representations. Moreover, recognizing that the classification task and mask prediction task involve distinct features, we choose to decouple these tasks. This strategic separation enables each query to focus on the most relevant features for its respective task. To address challenges stemming from partial label absence and domain shift, we incorporate pseudo-labels as an additional supervision signal. Experimental results obtained from partially supervised segmentation tasks on retinal fundus images validate the effectiveness of TAFormer, surpassing the performance of state-of-the-art methods.