Guided Frequency Loss (GFL) for Image Restoration

Official PyTorch implementation of Guided Frequency Loss (GFL), a novel loss function for image restoration tasks including Super-Resolution and Denoising.

📖 Abstract

Recent advancements in multimedia image enhancement and restoration have leveraged various generative models to achieve significant improvements. Despite this progress, the potential of utilizing the frequency domain remains underexplored, which is crucial for effectively addressing complex scenarios.

This repository introduces the Guided Frequency Loss (GFL), a novel approach designed to optimize the learning of frequency content in balance with the spatial patterns in images. It aggregates three major components that work in parallel to enhance learning efficiency:

• Charbonnier component
• Laplacian Pyramid component
• Gradual Frequency component

✨ Key Features

• 🎯 Joint Learning: Simultaneously learns spatial structures and frequency details for image restoration
• 🚀 Enhanced Training: Improves training stability and convergence across diverse deep models (Transformer & GAN-based)
• 📈 Progressive Strategy: Utilizes Gradual Frequency approach to progressively capture complex high-frequency components
• 🎨 Better Reconstruction: Strengthens texture and edge reconstruction through integrated loss terms
• 💪 Data Efficient: Improves restoration accuracy (PSNR/SSIM) even with limited training data

Tested environment

The current code was tested with:

• Linux
• Python 3.10.20
• PyTorch 2.2.2
• torchvision 0.17.2
• PyTorch Lightning 2.6.5
• timm 1.0.28

A CUDA GPU is optional for the smoke test. A GPU is strongly recommended for the full SwinIR configuration.

1. Clone the repository

git clone https://github.com/riotu-lab/guided-frequency-loss.git
cd guided-frequency-loss

Run all commands below from the repository root.

2. Create an environment

Recommended: Conda

conda create -n gfl python=3.10 -y
conda activate gfl
python -m pip install --upgrade pip

Install one PyTorch build. For CPU only:

python -m pip install torch==2.2.2 torchvision==0.17.2 \
    --index-url https://download.pytorch.org/whl/cpu

For an NVIDIA GPU with CUDA 12.1-compatible drivers:

python -m pip install torch==2.2.2 torchvision==0.17.2 \
    --index-url https://download.pytorch.org/whl/cu121

Then install this project and its training/test dependencies:

python -m pip install -e ".[training,dev]"

Verify the environment:

python -c "import torch; print('PyTorch:', torch.__version__); print('CUDA:', torch.cuda.is_available())"
python -m pytest tests/test_loss.py -q

The test command should report 23 passed.

Alternative: Python venv

python3.10 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e ".[training,dev]"
python -m pytest tests/test_loss.py -q

Use a terminal for installation and training commands. If a Jupyter kernel does not start, this does not prevent the terminal scripts from running.

3. Prepare the dataset

The super-resolution dataset expects clean HR RGB images only. The LR inputs are generated automatically by bicubic downsampling.

smoke_data/
├── train/
│   ├── image01.png
│   ├── image02.png
│   └── ...
└── test/
    ├── image01.png
    ├── image02.png
    └── ...

Supported extensions are PNG, JPEG, BMP, TIFF, and TIF. Subdirectories are searched recursively.

When val_dir: null, the training directory is split automatically according to train_split. Use at least two training images; four or more are recommended for a smoke test.

The current loader resizes every image to hr_size. For research training, square crops are usually preferable to stretching non-square source images.

4. Quick one-epoch smoke test

The quick configuration uses 32×32 HR targets, automatically creates 8×8 LR inputs, and trains a small 4× SwinIR for one epoch.

Train on CPU:

python scripts/train.py \
    --config configs/quick_smoke_sr_x4.yaml \
    --gpus 0

Train on one CUDA GPU:

python scripts/train.py \
    --config configs/quick_smoke_sr_x4.yaml \
    --gpus 1

Training writes:

checkpoints/quick_smoke_sr_x4/epoch-00.ckpt
checkpoints/quick_smoke_sr_x4/last.ckpt

Test the saved checkpoint on CPU:

python scripts/test.py \
    --config configs/quick_smoke_sr_x4.yaml \
    --checkpoint checkpoints/quick_smoke_sr_x4/last.ckpt \
    --gpus 0 \
    --batch-size 2

The test script reports PSNR and SSIM and writes comparison images to test_results/.

5. Train and test at 256×256

The provided configs/smoke_sr_x4.yaml configuration uses:

HR target:       256×256
LR model input:   64×64
scale:                 4×
window size:            8

Train:

python scripts/train.py \
    --config configs/smoke_sr_x4.yaml \
    --gpus 0

Use --gpus 1 for one CUDA GPU.

Test the final checkpoint:

python scripts/test.py \
    --config configs/smoke_sr_x4.yaml \
    --checkpoint checkpoints/smoke_sr_x4/last.ckpt \
    --gpus 0 \
    --batch-size 2

To test the best validation checkpoint, replace last.ckpt with the desired epoch-XX.ckpt file from the same checkpoint directory.

Do not reuse one checkpoint directory for configurations with different img_size, window_size, depths, embed_dim, or num_heads. Give each experiment a unique name, for example:

training:
  checkpoint_dir: checkpoints/sr_x4_256_w8
  experiment_name: sr_x4_256_w8

6. Image-size configuration rules

For 4× super-resolution, these values must agree:

model.img_size = data.hr_size / data.upscale_factor
model.upscale  = data.upscale_factor

The SwinIR window_size must divide the LR dimensions.

HR size	Scale	model.img_size	Valid example window	Adaptive start	Adaptive end
32	4	8	4	2	32
192	4	48	8	10	192
256	4	64	8	13	256

Example for 256×256 HR images:

model:
  upscale: 4
  img_size: 64
  window_size: 8

data:
  hr_size: [256, 256]
  upscale_factor: 4

If img_size is 64 but hr_size is 32, the loader creates an 8×8 LR image while the model expects 64×64. Training will then stop with AssertionError: input feature has wrong size.

7. Controlling Guided Frequency Loss

GFL is configured under the loss section:

loss:
  start_frequency: 13
  end_frequency: 256
  num_epoch: 100
  eps: 0.001

Parameters

Parameter	Meaning
start_frequency	Initial frequency-bin count and fixed low-column strip width.
end_frequency	Final row-strip cutoff reached at the end of the schedule.
num_epoch	Number of epochs over which the cutoff grows. Keep it equal to training.num_epochs.
eps	Charbonnier smoothing constant. 0.001 is the paper/default value.

8. Full SwinIR configuration

Edit the dataset paths in configs/swinir_sr_x4.yaml:

data:
  train_dir: /absolute/path/to/train/images
  val_dir: null
  test_dir: /absolute/path/to/test/images

Then run:

python scripts/train.py \
    --config configs/swinir_sr_x4.yaml \
    --gpus 1

The full configuration is much larger than the smoke model and is intended for GPU training.

9. Resume training

The resume checkpoint must use the same model configuration:

python scripts/train.py \
    --config configs/smoke_sr_x4.yaml \
    --gpus 1 \
    --resume checkpoints/smoke_sr_x4/last.ckpt

Project layout

guided-frequency-loss/
├── configs/                    # Quick, 256×256, and full configurations
├── gfl/
│   ├── data/                   # Dataset and transforms
│   ├── losses/                 # GFL and component losses
│   ├── models/                 # SwinIR
│   └── utils/                  # Metrics and visualization
├── notebooks/                  # Legacy interactive tutorial
├── scripts/
│   ├── train.py
│   ├── test.py
│   └── demo.py
├── smoke_data/                 # Small HR image dataset
├── tests/
├── requirements.txt
└── setup.py

Components

Class	Description
GradualFrequencyLoss	Complete spatial, Laplacian, and gradual-frequency loss.
CharbonnierLoss	Standalone smooth L1/Charbonnier penalty.
LaplacianPyramidLoss	Standalone depth-1 Laplacian-Pyramid loss.

Citation

If you use this work, please cite:

@article{benjdira2026gfl,
  title   = {Guided Frequency Loss for Multimedia Image Enhancement and Restoration},
  author  = {Benjdira, Bilel and Ali, Anas M. and Boulila, Wadii and Koubaa, Anis},
  journal = {Results in Engineering},
  year    = {2026},
  doi     = {10.1016/j.rineng.2026.111146}
}

License

MIT License. See LICENSE.

Acknowledgments

• SwinIR architecture: Liang et al., 2021
• Laplacian-Pyramid construction: Burt and Adelson, 1987