UnfoldIR: Rethinking Deep Unfolding Network in Illumination Degradation Image Restoration
Chunming He, Rihan Zhang, Fengyang Xiao, Chengyu Fang, Longxiang Tang, Yulun Zhang, Sina Farsiu
TL;DR
UnfoldIR addresses illumination degradation image restoration (IDIR) by treating image formation through Retinex as an optimization problem over reflectance $R$ and illumination $L$, with $I = R\odot L$ and objective $L(R,L) = \tfrac{1}{2}\|I - R\odot L\|_2^2 + \beta\varphi(R) + \gamma\phi(L) + \mu\mathcal{S}(\mathcal{A}(I) - \mathcal{A}(R)) + \tfrac{\lambda}{2}\| w\odot\nabla L\|_2^2$. The authors unfold the proximal-gradient iterations into UnfoldIR, a multistage network whose RAIC and IGRE modules leverage Visual State Space (VSS) concepts and a frequency-aware FVSS to refine illumination and texture, respectively, with a high-order RK2 solver and a learnable gating mechanism. A dedicated inter-stage information consistency loss $L_{ISIC}$ enforces stability across stages and enables strong unsupervised performance. Extensive experiments across 5 IDIR tasks and 3 downstream problems show state-of-the-art gains, supported by ablations and user studies, and the work provides practical guidance for designing future DUN-based restoration systems, including potential extensions with diffusion models and higher-level vision tasks.
Abstract
Deep unfolding networks (DUNs) are widely employed in illumination degradation image restoration (IDIR) to merge the interpretability of model-based approaches with the generalization of learning-based methods. However, the performance of DUN-based methods remains considerably inferior to that of state-of-the-art IDIR solvers. Our investigation indicates that this limitation does not stem from structural shortcomings of DUNs but rather from the limited exploration of the unfolding structure, particularly for (1) constructing task-specific restoration models, (2) integrating advanced network architectures, and (3) designing DUN-specific loss functions. To address these issues, we propose a novel DUN-based method, UnfoldIR, for IDIR tasks. UnfoldIR first introduces a new IDIR model with dedicated regularization terms for smoothing illumination and enhancing texture. We unfold the iterative optimized solution of this model into a multistage network, with each stage comprising a reflectance-assisted illumination correction (RAIC) module and an illumination-guided reflectance enhancement (IGRE) module. RAIC employs a visual state space (VSS) to extract non-local features, enforcing illumination smoothness, while IGRE introduces a frequency-aware VSS to globally align similar textures, enabling mildly degraded regions to guide the enhancement of details in more severely degraded areas. This suppresses noise while enhancing details. Furthermore, given the multistage structure, we propose an inter-stage information consistent loss to maintain network stability in the final stages. This loss contributes to structural preservation and sustains the model's performance even in unsupervised settings. Experiments verify our effectiveness across 5 IDIR tasks and 3 downstream problems.