Lightweight Temporal-frequency Perception Sparse State Space Models for Unified Image Restoration.

Abstract

Unified image restoration has become a fundamental issue in image processing. State space models have demonstrated significant potential in image restoration. However, their multi-directional scanning mechanism may introduce computational and feature redundancy, failing to satisfy lightweight deployment requirements. Furthermore, state space models have limitations in perceiving local detail features. To address this, we propose a lightweight channel-adaptive temporal-frequency sparse state space model for unified image restoration. This model enhances the local detail perception capability of the state space model using frequency domain features and simplifies the complexity of the network by sparse mechanisms. Specifically, we designed a U-shaped image restoration deep network based on the channel-adaptive temporal-frequency sparse state space module. This module consists of a temporal-domain dynamic sparse visual state space module and a frequency-domain sparse wavelet detail enhancement module in parallel, and uses a channel shuffling operation to realize temporal-frequency feature fusion. The dynamic sparse state space module uses a top-k mechanism to sparsify features across different scan paths for computational efficiency. The frequency-domain sparse wavelet detail enhancement module utilizes wavelet transformation and convolution operations to extract and enhance details in different directions, and then uses a top-k mechanism to perform sparse processing. Moreover, we introduce a degradation semantic perception module at the end of the encoder to guide the restoration network to adaptively learn the semantics of different degradation types, thereby realizing unified image restoration in complex outdoor environments. Extensive experimental results demonstrate that our method significantly outperforms 31 baseline methods in five complex weather and illumination degradation image restoration tasks while maintaining the lowest parameters and FLOPs.