TS-Diff: Two-Stage Diffusion Model for Low-Light RAW Image Enhancement

TL;DR

TS-Diff addresses extreme low-light RAW image enhancement where camera-specific noise distributions impede generalization. It introduces a two-stage diffusion framework that uses a noise-space of virtual cameras and Camera Feature Integration (CFI) to learn device-agnostic features, then forms a target-specific CFI^T via averaging and small real-data fine-tuning, with deployment aided by structural reparameterization. A Color Corrector (CC) stabilizes color distributions during diffusion, and a new Quantifiable Illumination Dataset (QID) enables controlled benchmarking across illumination levels. Experiments on SID, ELD, and QID demonstrate state-of-the-art denoising, cross-camera generalization, and color fidelity with efficient deployment. The work offers a practical, robust path toward real-world low-light RAW enhancement across diverse cameras.

Abstract

This paper presents a novel Two-Stage Diffusion Model (TS-Diff) for enhancing extremely low-light RAW images. In the pre-training stage, TS-Diff synthesizes noisy images by constructing multiple virtual cameras based on a noise space. Camera Feature Integration (CFI) modules are then designed to enable the model to learn generalizable features across diverse virtual cameras. During the aligning stage, CFIs are averaged to create a target-specific CFI$^T$, which is fine-tuned using a small amount of real RAW data to adapt to the noise characteristics of specific cameras. A structural reparameterization technique further simplifies CFI$^T$ for efficient deployment. To address color shifts during the diffusion process, a color corrector is introduced to ensure color consistency by dynamically adjusting global color distributions. Additionally, a novel dataset, QID, is constructed, featuring quantifiable illumination levels and a wide dynamic range, providing a comprehensive benchmark for training and evaluation under extreme low-light conditions. Experimental results demonstrate that TS-Diff achieves state-of-the-art performance on multiple datasets, including QID, SID, and ELD, excelling in denoising, generalization, and color consistency across various cameras and illumination levels. These findings highlight the robustness and versatility of TS-Diff, making it a practical solution for low-light imaging applications. Source codes and models are available at https://github.com/CircccleK/TS-Diff

Figures (8)

• Figure 1: Framework of the TS-Diff.
• Figure 2: Network architecture of the Color Corrector.
• Figure 3: Color Corrector in mitigating color shifts.
• Figure 4: Qualitative comparisons on SID dataset.
• Figure 5: Examples of images under varying illumination intensities. The first column displays the reference (ground truth) images, while the second, third, and fourth columns depict low-light images captured at illumination intensities of 10$^{-1}$ lux, 10$^{-2}$ lux, and 10$^{-3}$ lux, respectively.