| --- |
| library_name: pytorch |
| tags: |
| - image-denoising |
| - computer-vision |
| - u-net |
| - image-restoration |
| - cbam |
| - image-processing |
| metrics: |
| - psnr |
| - mae |
| - mse |
| datasets: |
| - BSDS500 |
| pipeline_tag: image-to-image |
| --- |
| |
| # LAR-U-Net: Lightweight Attentive Residual U-Net for Image Denoising |
|
|
|
|
| > **Course Project Notice:** This repository and model card were developed as part of the academic project for **`ARTI404: Image Processing`**. |
|
|
| **LAR-U-Net** is an advanced, lightweight image restoration architecture engineered specifically for grayscale image denoising. It structurally combines the spatial localization precision of a symmetric **U-Net** backbone, adaptive feature refinement via the **Convolutional Block Attention Module (CBAM)**, and a **Residual Noise Learning** paradigm. Additionally, it leverages a handcrafted **Gradient-Augmented Processing** front-end to strictly preserve structural boundaries and high-frequency edge dynamics during restoration. |
|
|
| --- |
|
|
| ## Architecture Overview |
|
|
| The primary challenge in discriminative image denoising is avoiding the over-smoothing of critical high-frequency details (e.g., fine textures, sharp edges). LAR-U-Net addresses this through a four-tiered architectural design: |
|
|
| 1. **Gradient-Augmented Processing (Front-End):** Instead of processing the raw noisy image in isolation, the input is passed through horizontal (G_x) and vertical (G_y) Sobel operators. The extracted gradient magnitude map (G) is concatenated along the channel dimension, forcing early spatial layers to prioritize structural edge maps. |
| 2. **Symmetric U-Net Backbone:** Utilizes an encoder-decoder architecture with skip connections, balancing parameter count and feature representation using optimized 3X3 convolutions. |
| 3. **Dual-Axis Attention Gates (CBAM):** Embedded directly inside each encoder stage to dynamically re-weight features: |
| * **Channel Attention:** Compresses spatial dimensions using adaptive global average pooling to capture cross-channel interdependencies (distinguishing signal from noise variants). |
| * **Spatial Attention:** Concatenates average and max pooling feature maps across the channel dimension and applies a 7X7 kernel to generate an explicit spatial mask highlighting destructive noise regions. |
| 4. **Residual Noise Learning Paradigm:** Adopting the discriminative mapping strategy popularized by **DnCNN**, the network directly predicts the underlying *Noise Map* R(y) rather than the clean image itself. |
|
|
| --- |
|
|
| ## Mathematical Formulation |
|
|
| ### 1. Image Degradation & Reconstruction Model |
| Given an observed noisy image y degraded by additive noise: |
|
|
| $$y = x + v$$ |
|
|
| Where x denotes the latent clean image and v represents the corrupting Additive White Gaussian Noise (AWGN) component with standard deviation. The restored image (x) is computed via explicit residual subtraction: |
|
|
| $$\hat{x} = y - R(y)$$ |
|
|
| ### 2. Front-End Edge Gradient Extraction |
| Spatial edge dynamics are computed using discrete Sobel spatial operators: |
|
|
| $$G = \sqrt{G_x^2 + G_y^2}$$ |
|
|
| ### 3. Optimization Criterion |
| The network parameters (W) are optimized by minimizing the Mean Squared Error (MSE) loss regularized explicitly over the target residual noise map: |
|
|
| $$\mathcal{L}(W) = \frac{1}{2N} \sum_{i=1}^{N} \| R(y_i; W) - (y_i - x_i) \|_F^2$$ |
| |
| --- |
| |
| ## Performance & Experimental Results |
| |
| ### Training Hyperparameters |
| * **Optimizer:** Adam |
| * **Loss Function:** Mean Squared Error (MSE) |
| * **Epochs:** 10 |
| * **Batch Size:** 8 |
| * **Target Noise Setup:** AWGN (sigma = 25) |
| * **Core Dataset:** Berkeley Segmentation Dataset 500 (**BSDS500**) — using disjoint train/val/test splits. |
| * **Robustness Evaluation Suite:** Multi Noises Dataset (Evaluated against Speckle, Poisson, Multiplicative, JPEG artifacts, Quantization, and Salt & Pepper noise). |
| |
| ### Quantitative Metrics |
| * **Minimum Training Loss:** `0.00315` (MSE) |
| * **Minimum Validation Loss:** `0.00327` (MSE) |
| * **Test Dataset MAE:** `0.04253` |
| * **Average Suite PSNR:** `19.9205 dB` |
| * **Peak PSNR Gain:** **> +7.0 dB** on structured natural scenes. |
| |
| #### Training Dynamics |
| | Epoch | Train Loss (MSE) | Validation Loss (MSE) | |
| | :---: | :--------------: | :------------------: | |
| | 1 | 0.22759 | 0.24682 | |
| | 2 | 0.00983 | 0.00924 | |
| | 5 | 0.00415 | 0.00414 | |
| | **10**| **0.00315** | **0.00327** | |
| |
| #### Sample-Level Reconstruction Analysis |
| | Evaluation Parameter | Sample 1 | Sample 2 | |
| | :------------------- | :------: | :------: | |
| | **Initial Noisy PSNR** | 20.21 dB | 20.19 dB | |
| | **Denoised PSNR** | 28.01 dB | 27.36 dB | |
| | **Net PSNR Gain** | **+7.79 dB** | **+7.18 dB** | |
| | **Initial MAE** | 0.0779 | 0.0783 | |
| | **Denoised MAE** | 0.0306 | 0.0331 | |
| | **Net MAE Reduction**| **-0.0473** | **-0.0452** | |
| |
| > **Architectural Note on Domain Generalization:** While LAR-U-Net yields quantitative improvements on impulsive noise types (e.g., Salt & Pepper), visual quality inspection confirms its performance is optimal for **Gaussian-distributed (AWGN) noise**. This is due to the inherent symmetry of residual noise mappings and Batch Normalization layers operating over normal probability distributions during SGD optimization. |
| |
| --- |
| |
| ## Comparison with Alternative Architectures & Trade-offs |
| |
| | Model Architecture | Key Strengths | Key Weaknesses | Optimization Loss | |
| | :--- | :--- | :--- | :--- | |
| | **Standard U-Net** | High spatial precision via skip connections; robust multi-scale localization. | High memory footprint; prone to producing blurry edges when trained on MSE alone. | MSE Loss / L1 Loss | |
| | **DnCNN** | Excellent Gaussian noise mapping via residual learning and Batch Normalization. | Lacks spatial attention; sensitive to structural domain shifts; computationally heavier at single resolution. | Residual MSE Loss | |
| | **Restormer** | State-of-the-art global contextual modeling via Transposed Self-Attention; exceptional benchmark results. | Extremely computationally expensive; heavy GPU memory overhead; requires huge datasets. | Frequency Loss | |
| | **LAR-U-Net** *(Ours)* | **Lightweight**; explicit edge preservation via Sobel front-end; spatial/channel feature refinement via **CBAM**. | Performance drops on heavy non-Gaussian noise distributions (e.g., extreme Salt & Pepper). | MSE Loss | |
