PyTorch Toolbox for Image Quality Assessment

An IQA toolbox with pure python and pytorch. Please refer to Awesome-Image-Quality-Assessment for a comprehensive survey of IQA methods and download links for IQA datasets.

📖 Introduction
⚡ Quick Start
- Dependencies and Installation
- Basic Usage
🥇 Benchmark Performances and Model Zoo
🛠️ Train
- Dataset Preparation
- Example Train Script

📖 Introduction

This is a image quality assessment toolbox with pure python and pytorch. We provide reimplementation of many mainstream full reference (FR) and no reference (NR) metrics (results are calibrated with official matlab scripts if exist). With GPU acceleration, most of our implementations are much faster than Matlab. Please refer to the following documents for details:

📦 Model Cards | 🗃️ Dataset Cards

🚩 Updates/Changelog

Oct 09, 2023. Add datasets: PIQ2023, GFIQA. Add metric topiq_nr-face. We release example results on FFHQ here for reference.
Aug 15, 2023. Add st-lpips and laion_aes. Refer to official repo at ShiftTolerant-LPIPS and improved-aesthetic-predictor
Aug 05, 2023. Add our work TOPIQ with remarkable performance on almost all benchmarks via efficient Resnet50 backbone. Use it with topiq_fr, topiq_nr, topiq_iaa for Full-Reference, No-Reference and Aesthetic assessment respectively.
March 30, 2023. Add URanker for IQA of under water images.
March 29, 2023. 🚨 Hot fix of NRQM & PI.
March 25, 2023. Add TreS, HyperIQA, CNNIQA, CLIPIQA.
More

⚡ Quick Start

Dependencies and Installation

Ubuntu >= 18.04
Python >= 3.8
PyTorch >= 1.12
Torchvision >= 0.13
CUDA >= 10.2 (if use GPU)

# Install with pip
pip install pyiqa

# Install latest github version
pip uninstall pyiqa # if have older version installed already 
pip install git+https://github.com/chaofengc/IQA-PyTorch.git

# Install with git clone
git clone https://github.com/chaofengc/IQA-PyTorch.git
cd IQA-PyTorch
pip install -r requirements.txt
python setup.py develop

Basic Usage

import pyiqa
import torch

# list all available metrics
print(pyiqa.list_models())

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# create metric with default setting
iqa_metric = pyiqa.create_metric('lpips', device=device)
# Note that gradient propagation is disabled by default. set as_loss=True to enable it as a loss function.
iqa_loss = pyiqa.create_metric('lpips', device=device, as_loss=True)

# create metric with custom setting
iqa_metric = pyiqa.create_metric('psnr', test_y_channel=True, color_space='ycbcr').to(device)

# check if lower better or higher better
print(iqa_metric.lower_better)

# example for iqa score inference
# Tensor inputs, img_tensor_x/y: (N, 3, H, W), RGB, 0 ~ 1
score_fr = iqa_metric(img_tensor_x, img_tensor_y)
score_nr = iqa_metric(img_tensor_x)

# img path as inputs.
score_fr = iqa_metric('./ResultsCalibra/dist_dir/I03.bmp', './ResultsCalibra/ref_dir/I03.bmp')

# For FID metric, use directory or precomputed statistics as inputs
# refer to clean-fid for more details: https://github.com/GaParmar/clean-fid
fid_metric = pyiqa.create_metric('fid')
score = fid_metric('./ResultsCalibra/dist_dir/', './ResultsCalibra/ref_dir')
score = fid_metric('./ResultsCalibra/dist_dir/', dataset_name="FFHQ", dataset_res=1024, dataset_split="trainval70k")

Example Test script

Example test script with input directory/images and reference directory/images.

# example for FR metric with dirs
python inference_iqa.py -m LPIPS[or lpips] -i ./ResultsCalibra/dist_dir[dist_img] -r ./ResultsCalibra/ref_dir[ref_img]

# example for NR metric with single image
python inference_iqa.py -m brisque -i ./ResultsCalibra/dist_dir/I03.bmp

🥇 Benchmark Performances and Model Zoo

Results Calibration

Please refer to the results calibration to verify the correctness of the python implementations compared with official scripts in matlab or python.

Performance Evaluation Protocol

We use official models for evaluation if available. Otherwise, we use the following settings to train and evaluate different models for simplicity and consistency:

Metric Type	Train	Test	Results
FR	KADID-10k	CSIQ, LIVE, TID2008, TID2013	FR benchmark results
NR	KonIQ-10k	LIVEC, KonIQ-10k (official split), TID2013, SPAQ	NR benchmark results
Aesthetic IQA	AVA	AVA (official split)	IAA benchmark results

Results are calculated with:

PLCC without any correction. Although test time value correction is common in IQA papers, we want to use the original value in our benchmark.
Full image single input. We use multi-patch testing only when it is necessary for the model to work.

Basically, we use the largest existing datasets for training, and cross dataset evaluation performance for fair comparison. The following models do not provide official weights, and are retrained by our scripts:

Metric Type	Reproduced Models
FR
NR	`cnniqa`, `dbcnn`, `hyperiqa`
Aesthetic IQA	`nima`, `nima-vgg16-ava`

Important Notes:

Due to optimized training process, performance of some retrained approaches may be different with original paper.
Results of all retrained models by ours are normalized to [0, 1] and change to higher better for convenience.
Results of KonIQ-10k, AVA are both tested with official split.
NIMA is only applicable to AVA dataset now. We use inception_resnet_v2 for default nima.
MUSIQ is not included in the IAA benchmark because we do not have train/split information of the official model.

Benchmark Performance with Provided Script

Here is an example script to get performance benchmark on different datasets:

# NOTE: this script will test ALL specified metrics on ALL specified datasets
# Test default metrics on default datasets
python benchmark_results.py -m psnr ssim -d csiq tid2013 tid2008

# Test with your own options
python benchmark_results.py -m psnr --data_opt options/example_benchmark_data_opts.yml

python benchmark_results.py --metric_opt options/example_benchmark_metric_opts.yml tid2013 tid2008

python benchmark_results.py --metric_opt options/example_benchmark_metric_opts.yml --data_opt options/example_benchmark_data_opts.yml

🛠️ Train

Dataset Preparation

You only need to unzip downloaded datasets from official website without any extra operation. And then make soft links of these dataset folder under datasets/ folder. Download links are provided in Awesome-Image-Quality-Assessment.
We provide common interface to load these datasets with the prepared meta information files and train/val/test split files, which can be downloaded from download_link and extract them to datasets/ folder.

You may also use the following commands:

mkdir datasets && cd datasets

# make soft links of your dataset
ln -sf your/dataset/path datasetname

# download meta info files and train split files
wget https://github.com/chaofengc/IQA-PyTorch/releases/download/v0.1-weights/meta_info.tgz
tar -xvf data_info_files.tgz

Examples to specific dataset options can be found in ./options/default_dataset_opt.yml. Details of the dataloader inferface and meta information files can be found in Dataset Preparation

Example Train Script

Example to train DBCNN on LIVEChallenge dataset

# train for single experiment
python pyiqa/train.py -opt options/train/DBCNN/train_DBCNN.yml

# train N splits for small datasets
python pyiqa/train_nsplits.py -opt options/train/DBCNN/train_DBCNN.yml

@misc{pyiqa,
  title={{IQA-PyTorch}: PyTorch Toolbox for Image Quality Assessment},
  author={Chaofeng Chen and Jiadi Mo},
  year={2022},
  howpublished = "[Online]. Available: \url{https://github.com/chaofengc/IQA-PyTorch}"
}

Please also consider to cite our new work TOPIQ if it is useful to you:

@misc{chen2023topiq,
      title={TOPIQ: A Top-down Approach from Semantics to Distortions for Image Quality Assessment}, 
      author={Chaofeng Chen and Jiadi Mo and Jingwen Hou and Haoning Wu and Liang Liao and Wenxiu Sun and Qiong Yan and Weisi Lin},
      year={2023},
      eprint={2308.03060},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}