██╗    ██╗ █████╗ ██╗   ██╗███████╗██╗     ███████╗████████╗
██║    ██║██╔══██╗██║   ██║██╔════╝██║     ██╔════╝╚══██╔══╝
██║ █╗ ██║███████║██║   ██║█████╗  ██║     █████╗     ██║   
██║███╗██║██╔══██║╚██╗ ██╔╝██╔══╝  ██║     ██╔══╝     ██║   
╚███╔███╔╝██║  ██║ ╚████╔╝ ███████╗███████╗███████╗   ██║   
 ╚══╝╚══╝ ╚═╝  ╚═╝  ╚═══╝  ╚══════╝╚══════╝╚══════╝   ╚═╝   
                                                              
     D E H A Z I N G   ·   F R E Q U E N C Y - A W A R E

Wavelet-based Deep Image Dehazing — replacing standard convolutions with Discrete Wavelet Transform layers to preserve high-frequency edge details during restoration.

---

✦ The Core Idea

Haze is a low-frequency phenomenon — it smoothly corrupts the global luminance of a scene. But scene content (edges, textures, object boundaries) lives in high frequencies. Standard convolutions treat all frequency bands identically, leading to blurred, over-smoothed restorations.

We fix this by inserting Discrete Wavelet Transform (DWT) layers at key processing stages. The DWT explicitly decomposes feature maps into sub-bands:

Input Feature Map X
        │
       DWT  (Haar, J=1 or J=2)
        │
   ┌────┴────┐
  LL        HF (LH, HL, HH)
  │          │
  ▼          ▼
Heavy      Lightweight
ConvBlock  HFBlock + Edge Attention
  │          │
  └────┬────┘
      IDWT
        │
  Restored Features

LL band → global haze structure → heavy processing
HF bands → edges & textures → lightweight processing + edge-preserving attention

---

✦ Results

Trained on RESIDE ITS — 1,000 training images, 200 validation images, 150 epochs, NVIDIA T4 GPU.

Model	Best PSNR ↑	Best SSIM ↑	Δ PSNR	Min Val Loss ↓
UNet Baseline	23.51 dB	0.9088	-0.46 dB	0.0682
WaveletUNetV1	23.89 dB	0.9100	-0.08dB	0.0661
WaveletUNetV2	24.77 dB	0.9296	+0.80dB	0.0606
FFA Baseline	23.97 dB	0.9104	—	0.0735
WaveletFFA V1	25.41 dB	0.9277	+1.44 dB	0.0586
WaveletFFA V2	24.63 dB	0.9247	+0.66 dB	0.0651
WaveletFFA V3	25.22 dB	0.9349	+1.25 dB	0.0586
WaveletFFA V3 (Hybrid)	25.53 dB	0.9360	+1.56 dB	0.0598

🏆 Every single wavelet-augmented model beats the baseline in both PSNR and SSIM.

---

✦ Model Zoo

We designed 7 models across two architectural families:

〉UNet Family

UNet Baseline — the convolution-only reference

Input (3ch)
  → Encoder: ConvBlock → MaxPool → ConvBlock → MaxPool
  → Bottleneck: ConvBlock (128ch)
  → Decoder: ConvTranspose + Skip Concat × 2
  → Output (3ch)

A clean two-level encoder-decoder with no frequency awareness. Serves as the ablation starting point.

WaveletUNet V1 — first wavelet bottleneck

Drops the standard bottleneck ConvBlock and replaces it with a WaveletBottleneck:

• Single-level Haar DWT on bottleneck features
• ConvBlock processes LL only — HF bands pass through unchanged
• IDWT reconstructs spatial features

yl, yh = self.dwt(x)         # decompose
ll = self.conv_ll(yl)        # process only low-freq
x  = self.idwt((ll, yh))     # reconstruct with original HF intact

WaveletUNet V2 — multi-band, multi-level, global residual

Deeper encoder (3 levels, up to 256ch) with WaveletBottleneckV2:

• 2-level Haar DWT (J=2)
• LL band: heavy ConvBlock + residual connection
• Level-1 HF bands (LH, HL, HH): each processed by a separate lightweight HFBlock (Conv-ReLU) + residual
• Level-2 HF bands: pass through unchanged
• Global input residual at output: clamp(output + input, 0, 1)

This shifts the learning target from full reconstruction to residual haze removal — a far simpler task.

# WaveletBottleneckV2 forward
yl, yh = self.dwt(x)                            # J=2 DWT
ll = self.conv_ll(yl) + yl                       # heavy LL + residual
# Level-1 HF: lightweight per-orientation blocks
lh = self.conv_lv1['lh'](level[:,:,0]) + level[:,:,0]
hl = self.conv_lv1['hl'](level[:,:,1]) + level[:,:,1]
hh = self.conv_lv1['hh'](level[:,:,2]) + level[:,:,2]
# Level-2 HF: pass through
return self.idwt((ll, new_yh))

---

〉FFA Family

All FFA variants share the lightweight backbone: channel dim = 32, 3 Groups × 8 Feature Attention Blocks = 24 total residual blocks. Each FAB applies Channel Attention (CA) + Pixel Attention (PA).

FFA Baseline — attention without wavelets

Standard FFA-Net architecture, lightly downsized for T4 GPU training:

• Original: 64ch, 19 FABs/Group → Ours: 32ch, 8 FABs/Group
• Reduces parameter count and prevents overfitting on 1,000 training samples
• Preserves all CA + PA attention mechanisms and adaptive multi-group fusion

Pre-conv (3→32ch)
  → Group 1: 8×FAB → Conv
  → Group 2: 8×FAB → Conv
  → Group 3: 8×FAB → Conv
  → Global CA fusion of [G1, G2, G3]
  → PA → Post-conv (32→3ch)
  → +input residual

WaveletFFA V1 — wavelet bottleneck after attention

Appends a WaveletBottleneck at the end of each Group, before the Group's residual addition:

Group(x) = WaveletBottleneck(Conv(8×FAB(x))) + x

Inside the WaveletBottleneck:

• Single-level Haar DWT on Group output features
• LL band: Conv-ReLU
• HF bands: orientation mean → PALayer edge attention → selective amplification of edge-containing sub-bands
• IDWT reconstruction

attn = self.edge_attention(yh[0].mean(dim=2))   # [B, C, H, W]
yh[0] = yh[0] * attn.unsqueeze(2)               # gate all 3 HF orientations
return self.idwt((yl, yh))

WaveletFFA V2 — Conv3D on HF bands + global shortcut

Upgrades WaveletFFA V1's HF processing from a simple attention gate to a 3D convolutional operator:

high_freq    = yh[0]                              # [B, C, 3, H/2, W/2]
processed    = self.edge_proc(high_freq)          # Conv3D treats orientations as depth
spatial_map  = processed.mean(dim=2)              # collapse orientations
attn         = self.edge_attention(spatial_map)   # PALayer
yh[0]        = processed * attn.unsqueeze(2)
out          = self.idwt((yl, yh)) + x            # global-local shortcut

Also uses a two-layer LL processor (Conv-ReLU-Conv) and a global residual wrapping the entire WaveletBottleneck.

WaveletFFA V3 — frequency FIRST, then attention

The most architecturally distinct variant. V1 and V2 place the wavelet module after the FABs. V3 inverts this order:

Group(x) = Conv(8×FAB(FrequencyBlock(x))) + x

The FrequencyBlock pre-processes inputs before attention:

• DWT → LL processed by Conv-ReLU-Conv
• HF bands processed by Conv3D
• Cross-band fusion: HF spatial map (via PALayer) modulates the LL branch before IDWT

hf_map  = self.hf_proc(hf).mean(dim=2)          # HF spatial summary
attn    = self.cross_attn(hf_map)                # cross-band attention
ll      = ll + attn                              # HF informs LL processing
out     = self.idwt((ll, yh)) + residual

By presenting frequency-separated, cross-band-refined inputs to the FABs, the attention modules can focus more precisely on informative spatial locations.

WaveletFFA V3 + Hybrid Loss — the best model 🏆

WaveletFFA V3 architecture, trained with the full hybrid loss combining four complementary objectives. Achieves 25.53 dB PSNR and 0.9360 SSIM — the top result across all experiments.

---

✦ Loss Function

All models support two loss modes:

Standard Loss

loss = L1(pred, target) + 0.2 × SSIM_Loss(pred, target)

Hybrid Loss (best results)

loss = L1 + 0.2×SSIM + 0.05×EdgeLoss + 0.05×FFTLoss

Component	What it does
L1	Absolute pixel accuracy, overall brightness
SSIM	Structural similarity — luminance, contrast, structure
Edge (Sobel)	Penalises differences in edge magnitude between pred and GT
FFT	Matches global frequency magnitude spectrum

The edge and FFT terms provide output-level frequency supervision that is complementary to the architectural frequency decomposition of DWT — both operating on different levels of abstraction.

---

✦ Quick Start

Installation

git clone https://github.com/Luv-valecha/Wavelet_Image_Dehazing.git
cd Wavelet_Image_Dehazing

# Using conda (recommended)
conda env create -f environment.yml
conda activate dehazing

# Or pip
pip install -r requirements.txt

Dataset Setup

# Download RESIDE ITS
python dataset_download.py

# Split into train/val
python dataset_split.py --train_size 1000 --val_size 200

Expected directory structure:

dataset/
├── train/
│   ├── hazy/     # 1000 hazy images
│   └── clear/    # 1000 clear images
└── test/
    ├── hazy/     # 200 validation images
    └── clear/

Training

# Train best model (WaveletFFA V3 + Hybrid Loss)
python train.py \
  --model_type waveletffav3 \
  --loss hybrid \
  --epochs 150 \
  --batch_size 8 \
  --lr 1e-4 \
  --run_name wavelet_ffa_v3_hybrid

# Train FFA baseline
python train.py --model_type ffa --loss l1_ssim

# Train WaveletUNet V2
python train.py --model_type wavelet_v2 --loss hybrid

# Resume from checkpoint
python train.py --model_type waveletffav3 --loss hybrid --resume

All available model types

--model_type	Architecture
baseline	UNet Baseline
wavelet	WaveletUNet V1
wavelet_v2	WaveletUNet V2
ffa	FFA Baseline
waveletffa	WaveletFFA V1
waveletffav2	WaveletFFA V2
waveletffav3	WaveletFFA V3 (best)

Inference

python test.py \
  --model_type waveletffav3 \
  --checkpoint outputs/best_waveletffav3_model.pth \
  --test_hazy dataset/test/hazy \
  --test_clear dataset/test/clear

---

✦ Project Structure

Wavelet_Image_Dehazing/
│
├── models/
│   ├── unet_baseline.py       # UNet Baseline
│   ├── unet_wavelet.py        # WaveletUNet V1
│   ├── unet_wavelet_v2.py     # WaveletUNet V2  ← multi-band, multi-level
│   ├── FFA.py                 # FFA Baseline
│   ├── waveletFFA.py          # WaveletFFA V1   ← wavelet after FABs
│   ├── waveletFFA_v2.py       # WaveletFFA V2   ← Conv3D on HF bands
│   └── waveletFFA_v3.py       # WaveletFFA V3   ← frequency-first ✦ BEST
│
├── losses/
│   ├── ssim_loss.py           # SSIM loss
│   ├── edge_loss.py           # Sobel edge loss
│   ├── fft_loss.py            # FFT magnitude loss
│   └── hybrid_loss.py         # Combined hybrid loss
│
├── utils/
│   ├── dataset_loader.py      # Training dataloader
│   ├── metrics.py             # PSNR & SSIM computation
│   └── visualization.py       # Validation visualisation
│
├── train.py                   # Training entry point
├── test.py                    # Evaluation entry point
├── dataset_download.py        # RESIDE ITS download script
├── dataset_split.py           # Train/val splitting
├── config.yaml                # Default configuration
└── environment.yml            # Conda environment

---

✦ Training Configuration

Hyperparameter	Value
Optimizer	Adam
Learning Rate	1e-4
Batch Size	8
Image Size	256 × 256
Epochs	150
Wavelet	Haar
FFA Channel Dim	32
FABs per Group	8 (24 total)
Hardware	NVIDIA T4 (16 GB)

---

✦ Why Wavelets?

The Haar wavelet transform is perfectly invertible, parameter-free, and decomposes signals into four interpretable sub-bands in a single pass:

         LL (avg × avg)    →  low-freq structure, haze
         LH (avg × diff)   →  horizontal edges
         HL (diff × avg)   →  vertical edges
         HH (diff × diff)  →  diagonal details, noise

For dehazing specifically:

• Haze lives in LL — it's a smooth, globally distributed degradation
• Content lives in HF — edges, textures, depth cues are high-frequency

Treating them separately with asymmetric computation — heavy for LL, lightweight-but-careful for HF — is a principled inductive bias that standard convolutions simply cannot provide.

---

✦ Acknowledgements

• RESIDE Dataset — Li et al., 2019
• FFA-Net — Qin et al., AAAI 2020
• pytorch-wavelets — Cotter, 2019
• Experiments tracked with Weights & Biases
• Compute: NVIDIA T4

---

"The Fourier transform tells you what frequencies are present.
The wavelet transform tells you what frequencies are present — and where."

Made with obsessive attention to scale · 2026