STUI-NET: Semi-Supervised Transformer for Underwater Information Enhancement

Underwater Image Restoration Technology (UIRT) constitutes a pivotal base for subsequent tasks yet confronts obstacles like data scarcity and image distortion in practical deployments. To optimize the use of scarce annotated and copious unlabeled data, we introduce an avant-garde semi-supervised method for underwater image restoration, named STUI-Net. Concurrently, we develop a novel teacher-student architecture utilizing Transformer technology, named GFT-Net. Initially, GFT-Net employs a tripartite branch network—comprising Gradient, Feature Extraction, and Transformer modules to thoroughly extract features from underwater images. This method then amalgamates multi-source features, encompassing edge, gradient, local, and global data, addressing underwater image deterioration in multifaceted environments. Additionally, we engineer a Supervisor role to rectify potential misdirection by the teacher network in the semi-supervised model, thereby enhancing training stability on unlabeled datasets. Empirical studies affirm the robust generalization competence of STUI-Net in authentic underwater environments.

1. INTRODUCTION

In recent years, with the rapid development of the underwater image restoration field [2], [6], [7], [11], [20], [26], researchers have paid extensive attention to key problems such as data scarcity and image distortion in this field. The unique nature of the underwater environment makes acquiring high-quality annotated data extremely challenging, thereby restricting the application of supervised learning methods in this domain. To address the problem of unlabeled data, a commonly adopted approach involves simulating underwater images using Generative Adversarial Networks [19]. However, this process often lacks stability, sometimes resulting in inconsistent quality of generated images or non-convergence of the loss function. Therefore, [4], [29] proposed a semi-supervised learning approach [24] to more effectively leverage limited annotated data and abundant unlabeled data, significantly enhancing the performance of underwater image restoration models.