IQA-PyTorch - 纯Python和PyTorch图像质量评估工具箱

PyTorch 图像质量评估工具箱

一个纯 Python 和 PyTorch 实现的 IQA 工具箱。请参阅 Awesome-Image-Quality-Assessment 获取全面的 IQA 方法调查和 IQA 数据集下载链接。

:open_book: 介绍

这是一个使用纯 Python 和 PyTorch实现的图像质量评估工具箱。我们重新实现了许多主流的全参考（FR）和无参考（NR）指标（如果存在官方的 Matlab 脚本，我们会进行结果校准）。在 GPU 加速下，我们的大部分实现比 Matlab 快得多。 详情请参阅以下文档：

📦 模型卡 | 🗃️ 数据集卡 | 🤗 数据集 | 📚 文档

:triangular_flag_on_post: 更新/变更日志

🔥2024年6月。添加 arniqa 及其在不同数据集上训练的各种变体，参见官方库这里。感谢 Lorenzo Agnolucci 的贡献 🤗。
2024年4月24日。添加 inception_score 和控制台入口命令 pyiqa。
2024年3月11日。添加 unique，参见官方库这里。感谢 Weixia Zhang 的贡献 🤗。
:boom: 2024年1月31日。添加了 NR 和 IAA 两种类型的 qalign。这是基于大型视觉语言模型的最强统一指标，性能和鲁棒性显著。更多详情请参考 Q-Align。使用代码如下：
```
qalign = create_metric('qalign').cuda()
quality_score = qalign(input, task_='quality')
aesthetic_score = qalign(input, task_='aesthetic')
```
2024年1月19日。添加 wadiqam_fr 和 wadiqam_nr。所有实现的方法现在都可用了 🍻。
2023年12月23日。添加 liqe 和 liqe_mix。感谢 Weixia Zhang 的贡献 🤗。
2023年10月9日。添加数据集：PIQ2023, GFIQA 。添加指标 topiq_nr-face。我们在 FFHQ 上发布了示例结果，参见这里。
更多

:zap: 快速开始

安装

# 使用 pip 安装
pip install pyiqa

# 安装最新的 github 版本
pip uninstall pyiqa # 如果已安装旧版本 
pip install git+https://github.com/chaofengc/IQA-PyTorch.git

# 使用 git clone 安装
git clone https://github.com/chaofengc/IQA-PyTorch.git
cd IQA-PyTorch
pip install -r requirements.txt
python setup.py develop

基本用法

你可以简单地使用命令行界面来使用此软件包。

# 列出所有可用的指标
pyiqa -ls

# 使用默认设置进行测试
pyiqa [metric_name(s)] --target [image_path or dir] --ref [image_path or dir]

代码高级用法

测试指标

import pyiqa
import torch

# 列出所有可用的指标
print(pyiqa.list_models())

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

# 使用默认设置创建指标
iqa_metric = pyiqa.create_metric('lpips', device=device)

# 检查是更低分数更好还是更高分数更好
print(iqa_metric.lower_better)

# IQA 评分推断示例
# 张量输入, img_tensor_x/y: (N, 3, H, W), RGB, 0 ~ 1
score_fr = iqa_metric(img_tensor_x, img_tensor_y)

# 图像路径作为输入
score_fr = iqa_metric('./ResultsCalibra/dist_dir/I03.bmp', './ResultsCalibra/ref_dir/I03.bmp')

# 对于 FID 指标，使用目录或预计算的统计数据作为输入
# 详情请参考 clean-fid: 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")

作为损失函数使用

请注意，默认情况下梯度传播是禁用的。设置 as_loss=True 以启用其作为损失函数使用。并非所有指标都支持反向传播，请参阅模型卡并确保以 lower_better 方式使用它。

lpips_loss = pyiqa.create_metric('lpips', device=device, as_loss=True)

ssim_loss = pyiqa.create_metric('ssimc', device=device, as_loss=True)
loss = 1 - ssim_loss(img_tensor_x, img_tensor_y)   # ssim 不是 lower better

使用自定义设置和权重

我们还提供了一种灵活的方式，以便您在重新训练或微调模型时使用自定义设置和权重。

iqa_metric = pyiqa.create_metric('topiq_nr', device=device, **custom_opts)

# 请注意，如果您使用此软件包训练模型，权重将保存在 weight_dict['params'] 中。否则，请设置 weight_keys=None。
iqa_metric.load_weights('path/to/weights.pth', weight_keys='params')

示例测试脚本

使用输入目录/图像和参考目录/图像的示例测试脚本。

# 目录方式的 FR 指标示例
python inference_iqa.py -m LPIPS[或 lpips] -i ./ResultsCalibra/dist_dir[dist_img] -r ./ResultsCalibra/ref_dir[ref_img]

# 单张图像的 NR 指标示例
python inference_iqa.py -m brisque -i ./ResultsCalibra/dist_dir/I03.bmp

:1st_place_medal: 基准性能和模型库

结果校准

请参阅结果校准以验证 Python 实现的正确性，这些实现与官方 Matlab 或 Python 脚本进行了比较。

⏬ 下载基准数据集

为了方便起见，我们将所有相关数据集上传到了 huggingface。以下是从 huggingface 下载数据集的示例代码：

[!注意] 我们仅收集这些数据集用于学术、研究和教育目的。用户应遵守每个数据集的原始创建者或所有者设定的使用指南、许可条款和条件。

import os
from huggingface_hub import snapshot_download

save_dir = './datasets'
os.makedirs(save_dir, exist_ok=True)

filename = "meta_info.tgz"
snapshot_download("chaofengc/IQA-Toolbox-Datasets", repo_type="dataset", local_dir=save_dir, allow_patterns=filename, local_dir_use_symlinks=False)

os.system(f"tar -xzvf {save_dir}/{filename} -C {save_dir}")

有关具体数据集选项的示例可以在 ./options/default_dataset_opt.yml 中找到。有关数据加载器接口和元信息文件的详细信息，请参阅数据集准备

性能评估协议

我们使用官方模型进行评估（如果有的话）。 否则，为了简化和一致性，我们使用以下设置来训练和评估不同的模型：

测量类型	训练	测试	结果
FR	KADID-10k	CSIQ, LIVE, TID2008, TID2013	FR基准结果
NR	KonIQ-10k	LIVEC, KonIQ-10k (官方划分), TID2013, SPAQ	NR基准结果
美学 IQA	AVA	AVA (官方划分)	IAA基准结果
结果计算如下：

PLCC 无任何修正。尽管在 IQA 论文中，测试时值修正是常见的做法，但我们希望在我们的基准测试中使用原始值。
全图单输入。我们不使用多块测试，除非必要。

基本上，我们使用现有最大的数据集进行训练，并使用跨数据集评估性能以进行公平比较。以下模型没有提供官方权重，是通过我们的脚本重新训练的：

指标类型	复现的模型
FR	`wadiqam_fr`
NR	`cnniqa`, `dbcnn`, `hyperiqa`, `wadiqam_nr`
美学 IQA	`nima`, `nima-vgg16-ava`

[!NOTE]

由于优化的训练过程，某些重新训练的方法性能可能与原始论文不同。

由我们重新训练的所有模型的结果都被归一化到 [0, 1]，且为了方便都改为了越高越好。

KonIQ-10k 和 AVA 的结果均使用官方划分进行测试。

NIMA 目前只适用于 AVA 数据集。我们对默认的 nima 使用 inception_resnet_v2。

MUSIQ 不包括在 IAA 基准测试中，因为我们没有官方模型的训练/划分信息。

使用提供的脚本进行基准性能测试

以下是一个在不同数据集上进行性能基准测试的示例脚本：

# 注意：此脚本将在所有指定的数据集上测试所有指定的指标
# 使用默认指标在默认数据集上进行测试
python benchmark_results.py -m psnr ssim -d csiq tid2013 tid2008

# 使用自定义选项进行测试
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

:hammer_and_wrench: 训练

训练脚本示例

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

# 单次实验训练
python pyiqa/train.py -opt options/train/DBCNN/train_DBCNN.yml

# 小数据集的 N 次划分训练
python pyiqa/train_nsplits.py -opt options/train/DBCNN/train_DBCNN.yml

分布式训练示例

torchrun --nproc_per_node=2 --master_port=4321 pyiqa/train.py -opt options/train/CLIPIQA/train_CLIPIQA_koniq10k.yml --launcher pytorch

:beers: 贡献

非常感谢对该代码库的任何贡献。请遵循贡献说明进行贡献指导。

:scroll: 许可

本作品采用 NTU S-Lab 许可证和知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议授权。

:bookmark_tabs: 引用

如果您认为我们的代码对您的研究有帮助，请考虑使用以下引用：

@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}"
}

如果我们的图像质量评估工作对您有帮助，也请考虑引用：

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

@article{wu2023qalign,
  title={Q-Align: Teaching LMMs for Visual Scoring via Discrete Text-Defined Levels},
  author={Wu, Haoning and Zhang, Zicheng and Zhang, Weixia and Chen, Chaofeng and Li, Chunyi and Liao, Liang and Wang, Annan and Zhang, Erli and Sun, Wenxiu and Yan, Qiong and Min, Xiongkuo and Zhai, Guangtai and Lin, Weisi},
  journal={arXiv preprint arXiv:2312.17090},
  year={2023},
  institution={Nanyang Technological University and Shanghai Jiao Tong University and Sensetime Research},
  note={Equal Contribution by Wu, Haoning and Zhang, Zicheng. Project Lead by Wu, Haoning. Corresponding Authors: Zhai, Guangtai and Lin, Weisi.}
}