README.md

metadata

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+vy\; =\; 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:

$$\widehat{x}=y-R(y)\backslash 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}G\; =\; \backslash 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\mathrm{\parallel }R(y_i;W)-(y_i-x_i){\mathrm{\parallel }}_F^2\backslash mathcal\{L\}(W)\; =\; \backslash frac\{1\}\{2N\}\; \backslash sum\_\{i=1\}^\{N\}\; \backslash |\; R(y\_i;\; W)\; -\; (y\_i\; -\; x\_i)\; \backslash |\_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