用于图像质量评估的 PyTorch 工具箱

一个带有纯 python 和 pytorch 的 IQA 工具箱。请参阅Awesome-Image-Quality-Assessment以获得对 IQA 方法的全面调查，以及 IQA 数据集的下载链接。

:open_book: 简介

这是一个纯python和pytorch的图像质量评估工具箱。我们提供以下功能：

• :sparkles:综合。支持许多主流的全参考（FR）和无参考（NR）指标
• :sparkles:准确。使用官方 matlab 脚本（如果存在）对我们的实现进行结果校准。
• :sparkles:快。借助 GPU 加速，我们的大部分实现都比 Matlab 快得多。
• :sparkles:灵活。支持使用多个公共 IQA 数据集训练新的 DNN 模型
• :sparkles:可区分的。大多数方法都支持 pytorch 向后

以下是该项目中支持的方法和数据集的详细信息。

:triangular_flag_on_post: 更新/变更日志

• 2022 年 6 月 3 日。添加 FID 指标。有关详细信息，请参阅clean-fid 。
• 2022 年 3 月 11 日。添加预训练的 DBCNN、NIMA 和 PieAPP 的官方模型 paq2piq。
• 更多的

:hourglass_flowing_sand: TODO 列表

• :white_large_square：在不同的数据集上添加预训练模型。

:zap: 快速入门

依赖和安装

• Ubuntu >= 18.04
• 蟒蛇> = 3.8
• Pytorch >= 1.8.1
• CUDA >= 10.1（如果使用 GPU）

# Install with pip
pip install pyiqa

# Install latest github version
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

基本用法

import pyiqa
import torch

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

# create metric with default setting
iqa_metric = pyiqa.create_metric('lpips', device=torch.device('cuda'))
# 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=torch.device('cuda'), 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")

示例测试脚本

带有输入目录和参考目录的示例测试脚本。-i和-r选项也支持单个图像。

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

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

:hammer_and_wrench: 火车

数据集准备

• 您只需从官网解压下载的数据集，无需任何额外操作。然后对文件夹下的这些数据集文件夹进行软链接datasets/。Awesome-Image-Quality-Assessment中提供了下载链接。
• 我们提供通用接口来加载这些数据集，其中包含准备好的元信息文件和训练/验证/测试拆分文件，可以从download_link 下载并将它们提取到datasets/文件夹中。

您还可以使用以下命令：

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/data_info_files.tgz
tar -xvf data_info_files.tgz

可以在 中找到特定数据集选项的示例./options/default_dataset_opt.yml。数据加载器接口和元信息文件的详细信息可以在数据集准备中找到

示例火车脚本

在 LIVEChallenge 数据集上训练 DBCNN 的示例

# 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

:1st_place_medal：基准表演和模型动物园

结果校准

请参考结果校准来验证python实现的正确性，并与matlab或python中的官方脚本进行比较。

经典指标的表现

这是一个示例脚本，用于在不同数据集上获得性能基准：

# 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

部分指标的基准表现请参考FR 基准结果和NR 基准结果。

深度学习模型的表现

我们在这里报告 PLCC/SRCC。

小数据集，拆分 1 的验证

大数据集性能

:beers: 贡献

非常感谢对此存储库的任何贡献。请按照贡献说明进行贡献指导。

:scroll: 许可证

本作品采用知识共享署名-非商业性-相同方式共享 4.0 国际许可协议进行许可。

:heart: 致谢

代码架构是从BasicSR借来的。几个实现取自：IQA-optimization、Image-Quality-Assessment-Toolbox、piq、piqa、clean-fid

我们还要感谢以下公共存储库：MUSIQ、DBCNN、NIMA、HyperIQA、CNNIQA、WaDIQaM、PieAPP、paq2piq、MANIQA

:e-mail: 联系方式

如果您有任何问题，请发送电子邮件chaofenghust@gmail.com