A Multi-Task Learning and Knowledge Selection Strategy for Environment-Induced Color-Distorted Image Restoration

Abstract

Existing methods for restoring color-distorted images in specific environments typically focus on a singular type of distortion, making it challenging to generalize their application across various types of color-distorted images. If it were possible to leverage the intrinsic connections between different types of color-distorted images and coordinate their interactions during model training, it would simultaneously enhance generalization, address potential overfitting and underfitting issues during data fitting, and consequently lead to a positive performance boost. In this paper, our approach primarily addresses three distinct types of color-distorted images, namely dust-laden images, hazy images, and underwater images. By thoroughly exploiting the unique characteristics and interrelationships of these types, we achieve the objective of multitask processing. Within this endeavor, identifying appropriate correlations is pivotal. To this end, we propose a knowledge selection and allocation strategy that optimally distributes the features and correlations acquired by the network from the images to different tasks, enabling a more refined task differentiation. Moreover, given the challenge of difficult dataset pairing, we employ unsupervised learning techniques and introduce novel Transformer blocks, feedforward networks, and hybrid modules to enhance context relevance. Through extensive experimentation, we demonstrate that our proposed method significantly enhances the performance of color-distorted image restoration.