CLE Diffusion: Controllable Light Enhancement Diffusion Model

Many image priors have found their use for traditional single-image low-light enhancement. Some approaches implement local and global histogram equalization  (Pizer et al., 1990; Abdullah-Al-Wadud et al., 2007) to increase the contrast of the input image. Some other solutions  (Li et al., 2015; Zhang et al., 2012) consider the low-light images as inverted haze images and use dehazing methods on the inverted input image. Another popular line of work is based on the Retinex theory (Land, 1977), which separates the image into illumination and reflectance layers and performs simple transformations on top of them for the desired effects. SRIE (Fu et al., 2015) estimates both layers simultaneously using a weighted variational model, while LIME (Guo et al., 2017) refines the illumination layer estimate and uses the decomposed reflection layer as the final enhanced result. For noise suppression, JED (Ren et al., 2018) has made progress by utilizing sequential decomposition. However, hand-crafted models used in these methods require careful parameter tuning and have limited model capacity.

In recent years, various learning-based light enhancement methods have been introduced thanks to the rapid development of deep learning. LLNet and S-LLNet (Lore et al., 2017) utilize deep autoencoder-based approaches for contrast enhancement and denoising. They are trained using data synthesized with random Gamma corrections and Gaussian noise. Based on the Retinex theory, Retinex-Net (Chen et al., 2018b) assumes that paired aligned images share the same reflectance but have different illuminations. The authors collect a paired low-light and normal-light data set, which paves the way for the training of larger models. Since collecting well-aligned paired datasets (Liu et al., 2021; Bychkovsky et al., 2011) requires hard work, great effort has been put into searching for methods without the need for paired supervision. EnlightenGAN  (Jiang et al., 2021) utilizes a generative adversarial network framework to learn a powerful generator without using unpaired data. Zero-DCE (Guo et al., 2020) learns light enhancement through image-specific curve estimation. CERL  (Chen et al., 2022a) builds upon EnlightenGAN and incorporates plug-and-play noise suppression. In addition, many works investigate the performance of different architectures, including Recursive Band Network (Yang et al., 2020), Signal-to-Noise-Ratio (SNR) Prior-aware Transformer (Xu et al., 2022c), Multi-axis MLP (Tu et al., 2022), Normalizing Flow (Wang et al., 2021), and Half Wavelet Network (Fan et al., 2022). These existing methods generally enhance the images to a pre-defined brightness level learned from the training dataset and avoid the challenging real-world case of region controlling for light enhancement.

A controllable GAN network named ReCoRo (Xu et al., 2022b) is the closest baseline to our approach. It allows users to specify the areas and levels of enhancement. However, the authors assign the brightest illumination level for each image using its normal-light version from the training set. This inconsistency of brightness level across images unstabilizes network optimization and disturbs the user experience at inference time. ReCoRo (Xu et al., 2022b) works with imprecise region controls by baking in augmentations. Effective as it is for portrait images, this method requires re-training for novel object classes. In contrast, our CLE Diffusion framework produces visually pleasing results via a conditional diffusion framework. Our unified brightness level also allows for seamless and consistent brightness control. Moreover, our framework offers a user-friendly experience with region control accessible through a single click.

Denoising diffusion model (Sohl-Dickstein et al., 2015) is a deep generative model synthesizing data through an iterative denoising process. Diffusion models consist of a forward process that distorts clean images with noise and a reverse process that learns to reconstruct clean images. They have demonstrated outstanding image generation (Song and Ermon, 2019; Ho et al., 2020) capability with the help of various improvements in architecture design (Dhariwal and Nichol, 2021; Rombach et al., 2022), sampling guidance (Ho and Salimans, 2022), and inference cost (Salimans and Ho, 2022; Rombach et al., 2022; Vahdat et al., 2021). Equipped with large-scale image-pair datasets, many works scaled up the model  (Rombach et al., 2022; Nichol et al., 2021; Ramesh et al., 2022; Saharia et al., 2022b) to billions of parameters to work with the challenging text-to-image generation. Additionally, denoising diffusion models have succeeded in various high-level computer vision tasks, including 3D generation (Poole et al., 2022; Singer et al., 2023; Xu et al., 2022a; Gu et al., 2023; Watson et al., 2022), object detection (Chen et al., 2022b) and depth estimation (Saxena et al., 2023). Meanwhile, various image restoration tasks are also studied through diffusion models, including super resolution (Saharia et al., 2022c), image deblurring (Whang et al., 2022), adverse weather condition (Özdenizci and Legenstein, 2023), and image to image translation (Saharia et al., 2022a). Our CLE Diffusion lies in the conditional denoising diffusion model category and makes the first attempt to study its effectiveness in controllable light enhancement.

Diffusion model (Sohl-Dickstein et al., 2015) is a type of generative model that uses iterative refinement to generate data. The widely used DDPM model (Ho et al., 2020) consists of a forward and a reverse process. The forward process gradually adds Gaussian noise to a clean input image, formulated as $q(y_{t}|y_{t-1})=\mathcal{N}(y_{t};\sqrt{\beta_{t}}y_{t-1},(1-\beta_{t})I)$. Furthermore, the intermediate steps can be marginalized out to characterize the distribution into $q(y_{t}|y_{0})=\mathcal{N}(\sqrt{\bar{\alpha_{t}}}y_{0},(1-\bar{\alpha_{t}})I)$, where $\alpha_{t}=1-\beta_{t}$ and $\bar{\alpha_{t}}=\prod_{i=0}^{t}\alpha_{i}$. The sequence of $\beta_{t}$ is carefully designed so that the noisy images will converge to pure Gaussian random noise $\mathcal{N}(0,I)$ when $t$ reaches the end of the forward process.

In the reverse process of the DDPM model, a neural network is used to denoise the data samples. The process can be formulated as $p_{\theta}(y_{t-1}|y_{t})=\mathcal{N}(\hat{y_{0}},\frac{1-\bar{\alpha_{t-1}}}{1-\bar{\alpha_{t}}}\beta_{t}),$ where $\hat{y_{0}}=\frac{1}{\sqrt{\alpha_{t}}}(y_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(y_{t},t)).$ $\epsilon_{\theta}$ is implemented with a U-Net (Ronneberger et al., 2015) model to estimate the noise component from a noisy image. During inference time, we sample $y_{T}\sim\mathcal{N}(0,I)$ and gradually reduce the noise level until we reach a clean image $y_{0}$. To accelerate sampling, DDIM (Song et al., 2020) presents a deterministic sampling approach as follows,

Diffusion models are usually trained via optimizing the negative log-likelihood loss function, which is further simplified (Ho et al., 2020; Kingma et al., 2021) as follows:

Specifically for low-light enhancement, we need to generate coherent normal-light images $y$ that share the content with the input low-light image $x$. Instead of learning the one-to-one mapping between the two domains, we are interested in approximating the conditional distribution $p(y|x)$ using the available paired image samples from the dataset. Similar to existing attempts to utilize diffusion model for image restoration (Saharia et al., 2022a, c), we implement our CLE Diffusion by adapting the original DDPM (Ho et al., 2020) model to accept additional condition information. While the forward process remains as simple as distorting the image with some carefully attenuated Gaussian noise, the U-net in the reverse process takes $x$ as an additional input.

Empirically, we observe simple concatenation of low-light image $x$ and diffused element $y_{t}$ leads to unstable training for complex scenes. This is partially due to the intricate noise and the diverse lighting in low-light environments. To this end, we further parameterize the conditional diffusion process with two additional priors.

Our first motivation comes from the severe color distortion in low-light images. As discussed in earlier works (Xu et al., 2022b; Liu et al., 2021), unnatural color shifts are often observed when enhancing low-light images. To this end, we implement a color map to reduce color distortion by normalizing the range of three color channels in input images. Specifically, an input image $x$ can be decomposed into three channels:

where $x_{r}$ means the red channel of the image, $x_{g}$ means the green channel of the image, $x_{b}$ means the blue channel of the image. We then extract the maximum pixel value for the three channels respectively:

where for example, $x_{r,max}$ means the maximum value of pixels on the red channel. Overall, the color map can be formulated as follows:

Another main challenge for enhancing low-light images lies in the inevitable noise in low-light conditions. Numerous researchers (Ren et al., 2019; Li et al., 2018; Chen et al., 2018b) have made various attempts to model the noise formulation. More recently, Xu et al. (Xu et al., 2022c) adopted an SNR-aware transformer for effective low-light enhancement. Specifically, the SNR map is used to bring spatial attention to the low signal-to-noise-ratio region. The SNR map can be obtained as follows,

where $\epsilon$ is included for numerical stability and $F$ is a low-pass filter implemented as a Gaussian blur in our experiments. We consider the high-frequency component in the image to be noise and calculate the ratio between the original image and noise directly.

In each training step, we randomly sample a pair of low-light image $x$ and its corresponding normal-light image $y$. We then prepare the color map $C(x)$, the SNR map $S(x)$ and a noisy image $y_{t}=\sqrt{\bar{a}_{t}}{y}+\sqrt{1-\bar{a}_{t}}\epsilon$. They are concatenated with low-light image $x$ as the input of our diffusion model. Consequently, Eq. 2 can be extended as follows,

To emphasize the effectiveness of conditioning information, we adopt a classifier-free guidance (Ho and Salimans, 2022) approach. This involves jointly training a conditional and an unconditional diffusion model, without the need for a classifier on noisy images as in the classifier guidance technique (Dhariwal and Nichol, 2021). Instead, the classifier-free guidance approach allows for a trade-off between sample quality and diversity by adjusting the weighting of conditional sampling and unconditional sampling during inference. For CLE Diffusion, we treat the brightness level of images as our ”class” labels. It is important to note that our ”class” in this case is continuous, allowing for seamless interpolation and continuous adjustment of the target brightness levels.

We first extract the vanilla brightness level $\lambda$ of normal-light images by calculating the average pixel value. Then we encode the average value using a random orthogonal matrix into the illumination embedding. Specifically, we uniformly establish several discretized values in $[0,1]$ as anchors. Given a $n\times n$ encoding matrix, the $d$-th value is mapped to the $d$-th column of the $n$-dimensional random orthogonal matrix. We adopt bilinear interpolations of the two nearby columns for the intermediate values. The illumination embedding is further embedded into the U-net using our Brightness Control Module. As shown in Fig. 3, we adopt a FiLM layer (Perez et al., 2018) to learn feature-wise affine transformation based on the illumination embedding. Then the modulated feature is split by half along the channel axis, with one copy being multiplied by the feature map and the other added to the feature map.

We train the conditional diffusion model $\epsilon_{\theta}(y_{t},x,C(x),S(x)|\lambda)$ together with an unconditional model $\epsilon_{\theta}(y_{t},x,C(x),S(x)|0)$. To be more specific, for the unconditional model, we input a zero embedding with the same shape as the illumination embedding. During the training process, we randomly drop the conditioning by allowing $2\%$ of the total iterations to train the unconditional model.

As shown in Algo. 1, the sampling process is implemented with DDIM (Song et al., 2020) sampler. Firstly, we sample $y_{T}\sim\mathcal{N}(0,I)$ . Subsequently, we estimate two noise maps, one from the conditional model and the other from the unconditional model, and apply a weighted average of the two noise estimates. The guidance scale $s$ is used to regulate the influence of the conditioning signal, where a larger scale produces results more aligned with the controlling signal, while a smaller scale produces less connected results.

Users may prioritize increasing the brightness of specific regions of interest over globally illuminating the entire image, especially when dealing with complex lighting conditions. To address this need, we introduce region controllability to our CLE Diffusion method, referred to as Mask-CLE Diffusion.

We incorporate a binary mask $M$ into our diffusion model by concatenating the mask with the original inputs. To accommodate this requirement, we created synthetic training data by randomly sampling free-form masks (Suvorov et al., 2022) with feathered boundaries. The target images are generated by alpha blending the low-light and normal-light images from existing low-light datasets (Chen et al., 2018b; Bychkovsky et al., 2011).

Segment-Anything Model (SAM) (Kirillov et al., 2023) is a large vision transformer trained on 11 million images to produce high-quality object masks from input prompts such as points or boxes. With the aid of SAM (Kirillov et al., 2023), obtaining precise object masks in low-light conditions becomes a user-friendly process. Users can select their desired region with just one click. Mask-CLE Diffusion subsequently generates controllable light enhancement images via mixing the results from $\epsilon_{\theta}(y_{t},x,C(x),S(x),M|\lambda)$ and $\epsilon_{\theta}(y_{t},x,C(x),S(x),M|0)$. For tiny objects such as human hair, SAM masks are not accurate enough. In such cases, we can still use the SAM mask as an initialization for downstream models. In our experiments, we utilize MatteFormer (Park et al., 2022) to automatically construct detailed alpha mattes for human matting. The trimaps used for MatteFormer are constructed using dilate and erode operations.

Although an ideal optimum for the $\mathcal{L}_{\text{simple}}$ will be able to approximate the $p(y|x)$ distribution from the available paired dataset, in practice, we observe that the model often ends up generating color distortion and unprecedented noise. To improve the convergence of our diffusion model, we introduce auxiliary losses to provide direct supervision on the estimated denoised images $\hat{y_{0}}$.

Existing quantitative metrics typically assume the existence of an ideal brightness level, making it difficult to compare images with different brightness levels fairly. As shown in Fig. 4, PSNR and SSIM present an inverted “V”-shaped pattern while LPIPS show a “V”-shaped pattern. These metrics exhibit large variances as brightness levels change, indicating their sensitivity to changes in light intensity. As light enhancement is highly subjective, a better value of these metrics does not necessarily correlate with better image quality. Additionally, the highly ill-posed nature of light enhancement brings about many possible optimal solutions. These solutions can have different white balances and exposure levels, further complicating the determination of a single ”best” solution.

We are initially inspired by the color map’s ability to normalize the brightness of images. As mentioned in Eq. 5, the color map can bring images with different brightness levels to a normalized brightness level, especially for images that are too dark or overexposed. Then, we use Canny edge detector (Canny, 1986) to extract useful structural information, further reducing the appearance information that is easily affected by illumination. We have empirically set the Canny threshold as 50 and 250. Other combinations exhibit minimal variation. As semantic features are more resilient than pixel-space metrics, we select the LPIPS distance calculation for comparing the structure of two images. We name this new metric as Light-Independent LPIPS, dubbed $LI\text{-}LPIPS$:

where $C(\cdot)$ is color map, $\phi^{l}$ is the $l$-th layer feature map of a pre-trained VGG network, and $\text{Canny}(\cdot)$ represents the Canny operator. In Fig.4, when the average image brightness ranges from 0.18 to 0.69, the difference observed in LI-LPIPS is less than 0.02, while LPIPS yields a variation in maximum and minimum values as high as 0.15. This validates LI-LPIPS is insensitive to variations in brightness.

We evaluate the CLE-Diffusion on two popular benchmarks, LOL (Chen et al., 2018b) and MIT-Adobe FiveK (Bychkovsky et al., 2011). The LOL dataset contains 485 training and 15 testing paired images, with each pair comprising a low-light and a normal-light image. MIT-Adobe FiveK consists of 5000 images processed with Adobe Lightroom by five experts. Following the split in previous methods (Tu et al., 2022; Ni et al., 2020), 4500 paired images are utilized as the training set, while the remaining 500 paired images serve as the test set. We use PSNR, SSIM, LPIPS (Zhang et al., 2018a), and LI-LPIPS metrics to measure the quality of output images.

We compare our CLE Diffusion with state-of-the-art low-light image enhancement methods. For quantitative comparison, we select several representative methods, including deep Retinex-based methods (RetinexNet (Chen et al., 2018b), KinD++ (Zhang et al., 2019)), CNN-based methods (MAXIM (Tu et al., 2022), HWMNet (Fan et al., 2022)), a Zero-Shot learning method (Zero-DCE (Guo et al., 2020)), a Semi-Supervised Learning method (DRBN (Yang et al., 2020)), a GAN-based (EnlightenGAN (Jiang et al., 2021)) and a Flow-based model (LLFlow  (Wang et al., 2022)).

As shown in Tab. 1, our method surpasses the compared methods in terms of PSNR and ranks third in SSIM, LPIPS, and LI-LPIPS. These findings provide evidence that our CLE Diffusion technique can generate superior samples from the dataset distribution in comparison to current approaches.

Visual comparisons are shown in Fig. 5, our proposed method surpasses other methods by accurately aligning the predicted luminance with the ground truth. The outputs generated by our framework exhibit significantly higher quality, characterized by well-suppressed noise, in contrast to the results produced by other methods. These alternative approaches display color shifts, noisy artifacts, and irregular illumination.

To further validate the ability of CLE Diffusion, we test on the MIT-Adobe FiveK dataset, which is much larger than LOL dataset and includes more diverse scenes. Tab. 2 shows the comparison results with other methods. Our method achieves the highest values in terms of PSNR, while also producing results comparable to state-of-the-art methods in terms of SSIM. As indicated by Fig. 6, our method is more consistent with ground truth in terms of color distortion. Compared to previous methods, our methods can preserve better details and color consistency as exemplified in zoomed-in regions.

With the powerful capabilities of CLE-Diffusion and Mask-CLE-Diffusion, we can not only control the global brightness of images at a specified brightness level but also achieve precise brightness control for desired regions. Users can input multiple desired brightness levels into the network to sample multiple images with various brightness levels as depicted in Fig. 7. For images with distinct brightness levels, the fidelity of details is remarkably high and the images exhibit no apparent color distortion, overexposure, or underexposure while maintaining a consistent global brightness.

By utilizing SAM model (Kirillov et al., 2023), users can effortlessly obtain regions of interest(ROI) by clicking on the image. As demonstrated in Fig. A, Mask-CLE-Diffusion is able to selectively enhance specific regions to attain desired brightness levels. In contrast to directly blending normal-light and low-light images using masks, our results have a more natural appearance and better emphasize the ROI.

We conduct ablation studies to validate the effectiveness of our proposed loss functions and network design. As shown in Fig. 8, when training with only $\mathcal{L}_{\text{simple}}$, the output images suffer from color distortion, artifacts, and unsatisfactory lighting. And when training without the conditioning color map and SNR map, the network produces noisy results which do not appear in the results of the full model. Tab. 3 also presents a notable decrease in all four metrics. These results show that auxiliary losses, the conditioning color map, and the SNR map contribute to improving the overall performance of the model.