PMRF 是一种新颖的逼真图像修复算法。它（在理论上）近似于在完美感知质量约束下，能够最小化均方误差（MSE）的最优估计器。

---

📈 我们论文中的部分结果

CelebA-Test 定量对比

红色、蓝色和绿色分别表示最佳、第二佳和第三佳的成绩。

WIDER-Test 视觉对比

WebPhoto-Test 视觉对比

⚙️ 安装

Windows 用户注意： 似乎有几位 Windows 用户无法安装 natten 包，而该包是使用 PMRF 中 HDiT 模型架构所必需的。一个对多人有效的解决方案已在 这里 提出。如果您仍未能解决此问题，可以尝试使用其他架构（如 UNet）来训练 PMRF，从而避免使用 natten 包。

我们通过以下命令创建了一个 conda 环境，务必严格按照给定顺序执行（这些命令也包含在 install.sh 文件中）：

conda create -n pmrf python=3.10
conda activate pmrf
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=11.8 -c pytorch -c nvidia
conda install lightning==2.3.3 -c conda-forge
pip install opencv-python==4.10.0.84 timm==1.0.8 wandb==0.17.5 lovely-tensors==0.1.16 torch-fidelity==0.3.0 einops==0.8.0 dctorch==0.1.2 torch-ema==0.3
pip install natten==0.17.1+torch230cu118 -f https://shi-labs.com/natten/wheels
pip install nvidia-cuda-nvcc-cu11
pip install basicsr==1.4.2
pip install git+https://github.com/toshas/torch-fidelity.git
pip install lpips==0.1.4
pip install piq==0.8.0
pip install huggingface_hub==0.24.5

1. 请注意，PMRF 使用的 HDiT 架构需要 natten 包。请确保将 natten==0.17.1+torch230cu118 替换为您系统上安装的正确 CUDA 版本。可用版本请参阅 https://shi-labs.com/natten/。
2. 我们安装了 nvidia-cuda-nvcc-cu11，因为否则 torch.compile 会无故卡住。不过，在您的系统上，torch.compile 可能无需此包即可正常运行。无论如何，如果您希望如此，可以直接跳过该包，或从我们的代码中移除所有涉及 torch.compile 的行。
3. 由于 basicsr 存在兼容性问题，您需要修改该包中的一个文件。打开 /path/to/env/pmrf/lib/python3.10/site-packages/basicsr/data/degradations.py，其中 /path/to/env 是您的 conda 环境安装路径。 然后将以下行：

from torchvision.transforms.functional_tensor import rgb_to_grayscale

改为：

from torchvision.transforms.functional import rgb_to_grayscale

⬇️ 下载

🌐 模型检查点

我们提供了盲脸图像修复模型的检查点，可在 Hugging Face 和 Google Drive 上获取。论文第 5.2 节中受控实验的检查点也可从 Google Drive 下载。请保持与 Google Drive 中相同的文件夹结构：

checkpoints/
├── blind_face_restoration_pmrf.ckpt    # 盲脸图像修复模型检查点。
├── swinir_restoration512_L1.pth    # DifFace 训练的 SwinIR 模型检查点。
├── controlled_experiments/     # 受控实验的检查点
│   ├── colorization_gaussian_noise_025/
│   │   ├── pmrf/
│   │   │   └── epoch=999-step=273000.ckpt
│   │   ├── mmse/
│   │   │   └── epoch=999-step=273000.ckpt
.   .   .
.   .   .
.   .   .

为了评估论文中的地标距离（LMD）和身份度量（Deg），您还需要从 VQFR 的 Google Drive 下载 resnet18_110.pth 和 alignment_WFLW_4HG.pth 检查点，并将其放置在 evaluation/metrics_ckpt/ 文件夹中。

🌐 盲脸图像修复测试数据集

1. 从 https://xinntao.github.io/projects/gfpgan 下载 WebPhoto-Test、LFW-Test 和 CelebA-Test（HQ 和 LQ）。
2. 从 https://shangchenzhou.com/projects/CodeFormer/ 下载 WIDER-Test。
3. 将这些数据集存放在您系统中的任意位置。

🧑 盲脸图像修复（论文第 5.1 节）

⚡ 快速推理

为了快速使用我们的模型，我们提供了一个 Hugging Face 检查点，它会自动下载。只需运行以下命令：

python inference.py \
--ckpt_path ohayonguy/PMRF_blind_face_image_restoration \
--ckpt_path_is_huggingface \
--lq_data_path /path/to/lq/images \
--output_dir /path/to/results/dir \
--batch_size 64 \
--num_flow_steps 25

请根据需要调整 --num_flow_steps（这是我们论文中的超参数 K）。

你也可以提供本地的模型检查点（例如，如果你训练了自己的 PMRF 模型，或者希望使用我们的 Google Drive 检查点而不是 Hugging Face 的）。只需运行以下命令：

python inference.py \
--ckpt_path ./checkpoints/blind_face_restoration_pmrf.ckpt \
--lq_data_path /path/to/lq/images \
--output_dir /path/to/results/dir \
--batch_size 64 \
--num_flow_steps 25

需要注意的是，我们的盲人脸图像修复模型是针对方形且对齐的人脸图像进行训练的。如果要修复包含多个人脸的一般内容图像，可以使用我们的 Hugging Face 演示。

🔬 评估

1. 我们从 VQFR 的 Google Drive 下载了 resnet18_110.pth 和 alignment_WFLW_4HG.pth 检查点，并将它们放入 evaluation/metrics_ckpt/ 文件夹中。 要在 CelebA-Test 数据集上评估结果，请运行：

   cd evaluation
   python compute_metrics_blind.py \
   --parent_ffhq_512_path /path/to/parent/of/ffhq512 \
   --rec_path /path/to/celeba-512-test/restored/images \
   --gt_path /path/to/celeba-512-test/ground-truth/images
   

   要在真实世界数据集上评估结果，请运行：

   cd evaluation
   python compute_metrics_blind.py \
   --parent_ffhq_512_path /path/to/parent/of/ffhq512 \
   --rec_path /path/to/real-world/restored/images \
   --mmse_rec_path /path/to/mmse/restored/images
   

   其中 --mmse_rec_path 参数是可选的，用于计算 IndRMSE，作为真实世界退化图像的真实 RMSE 指标。请注意，当你运行 inference.py 时，MMSE 重建结果会自动保存，因为 MMSE 模型也包含在 PMRF 检查点中。

💻 训练

在 scripts/ 文件夹中，我们提供了用于盲人脸图像修复以及基线模型训练的脚本。如果你想运行某个脚本，需要在根目录下执行（即 train.py 所在的目录）。训练该模型需要 FFHQ 数据集。我们下载了原始的 FFHQ 1024×1024 数据集，并使用双三次下采样将其缩小到 512×512。

1. 将位于 scripts/train/blind_face_restoration 中的 train_pmrf.sh 文件复制到根目录。
2. 根据你系统中训练和验证数据的位置，调整 --train_data_root 和 --val_data_root 参数。
3. 由 DifFace 训练的 SwinIR 模型已放在 checkpoints/ 文件夹中。我们通过以下命令下载了它：

   wget https://github.com/zsyOAOA/DifFace/releases/download/V1.0/swinir_restoration512_L1.pth
   

4. 将 --mmse_model_ckpt_path 参数调整为 SwinIR 模型的路径。
5. 根据你的系统配置，调整 --num_gpus 和 --num_workers 参数。
6. 运行 train_pmrf.sh 脚本以训练我们的模型。

👩‍🔬 对照实验（论文第 5.2 节）

我们提供了论文中对照实验的训练和评估代码，其中我们将 PMRF 与以下基线方法进行了比较：

1. 基于 Y 的流模型：一种经过修正的流模型，它以输入测量值为条件，学习从纯噪声流到真实数据分布。
2. 基于后验均值预测器的流模型：一种经过修正的流模型，它以后验均值预测为条件，学习从纯噪声流到真实数据分布。
3. 直接从 Y 流动的模型：一种经过修正的流模型，它从退化的测量值流到真实数据分布。
4. 后验均值预测器：一种旨在最小化 MSE 损失的模型。

🔬 评估

我们提供了用于快速评估 PMRF 及所有基线方法的检查点。

1. 评估是在尺寸为 256×256 的 CelebA-Test 图像上进行的。为了获取这些图像，我们从 GFPGAN 下载了 CelebA-Test (HQ) 图像，并使用双三次下采样将其缩小到 256×256。
2. 在 test.sh 中将 --test_data_root 调整为 CelebA-Test 256×256 图像的路径，并根据你要评估的退化类型及相应模型检查点，调整 --degradation 和 --ckpt_path 参数。
3. 运行 test.sh。

我们会自动保存重建后的输出、退化后的测量值以及来自源分布的样本（即 ODE 求解器开始处理的图像）。运行 test.sh 后，你可以通过以下方式评估结果：

cd evaluation
python compute_metrics_controlled_experiments.py \
--parent_ffhq_256_path /path/to/parent/of/ffhq256 \
--rec_path /path/to/restored/images \
--gt_path /path/to/celeba-256-test/ground-truth/images

💻 训练

• 我们在 FFHQ 256×256 数据上训练了模型。为了获得这些图像，我们将原始 FFHQ 1024×1024 图像使用双三次下采样缩小到了 256×256。
• PMRF 以及每个基线模型的训练脚本都位于 scripts/train/controlled_experiments/ 文件夹中。
• 要运行这些脚本，你需要将其复制到 train.py 所在的根目录。你只需调整每个脚本中的 --degradation、--source_noise_std、--train_data_root 和 --val_data_root 参数即可。对于去噪任务，我们使用了 --source_noise_std 0.025；而对于其他任务，则使用了 --source_noise_std 0.1。
• 要运行 train_pmrf.sh 和 train_posterior_conditioned_on_mmse_model.sh 脚本，首先需要通过 train_mmse.sh 训练 MMSE 模型。然后，根据 MMSE 模型最终检查点的路径，调整 --mmse_model_ckpt_path 参数。

📝 引用

@inproceedings{
    ohayon2025posteriormean,
    title={Posterior-Mean Rectified Flow: Towards Minimum {MSE} Photo-Realistic Image Restoration},
    author={Guy Ohayon and Tomer Michaeli and Michael Elad},
    booktitle={The Thirteenth International Conference on Learning Representations},
    year={2025},
    url={https://openreview.net/forum?id=hPOt3yUXii}
}

📋 许可证与致谢

本项目采用 MIT 许可证 开源。

我们借鉴了 BasicSR、VQFR、DifFace、k-diffusion 和 SwinIR 的代码。感谢这些仓库的作者提供了有用的实现。

📧 联系方式

如果您有任何问题或咨询，请随时联系我。

PMRF 快速上手指南

PMRF (Posterior-Mean Rectified Flow) 是一种新颖的照片级真实感图像修复算法，旨在在保持完美感知质量约束的同时，最小化均方误差 (MSE)。本指南将帮助您快速配置环境并运行盲人脸图像修复模型。

1. 环境准备

• 操作系统: 推荐 Linux (Ubuntu)。
  • Windows 用户注意: natten 包在 Windows 上安装可能失败。若遇到此问题，建议改用 UNet 架构训练，或参考官方 Issue #8 中的解决方案。
• 硬件要求: NVIDIA GPU (支持 CUDA 11.8)。
• 软件依赖:
  • Conda (推荐用于环境管理)
  • Python 3.10
  • CUDA Toolkit 11.8

2. 安装步骤

请严格按照以下顺序执行命令以创建名为 pmrf 的 Conda 环境并安装依赖。

2.1 创建环境与安装基础库

conda create -n pmrf python=3.10
conda activate pmrf
conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=11.8 -c pytorch -c nvidia
conda install lightning==2.3.3 -c conda-forge

2.2 安装 Python 依赖包

pip install opencv-python==4.10.0.84 timm==1.0.8 wandb==0.17.5 lovely-tensors==0.1.16 torch-fidelity==0.3.0 einops==0.8.0 dctorch==0.1.2 torch-ema==0.3
pip install nvidia-cuda-nvcc-cu11
pip install basicsr==1.4.2
pip install git+https://github.com/toshas/torch-fidelity.git
pip install lpips==0.1.4
pip install piq==0.8.0
pip install huggingface_hub==0.24.5

2.3 安装 NATTEN (关键步骤)

natten 是 PMRF 使用的 HDiT 架构所必需的。请务必根据您的实际 CUDA 版本替换下方命令中的版本号。以下示例适用于 CUDA 11.8 + PyTorch 2.3：

pip install natten==0.17.1+torch230cu118 -f https://shi-labs.com/natten/wheels

其他版本请参考：https://shi-labs.com/natten/

2.4 修复兼容性問題

由于 basicsr 包的兼容性问题，需要手动修改一个文件。 找到您的环境路径（例如 /path/to/env/pmrf/lib/python3.10/site-packages/basicsr/data/degradations.py），打开该文件，将：

from torchvision.transforms.functional_tensor import rgb_to_grayscale

修改为：

from torchvision.transforms.functional import rgb_to_grayscale

3. 基本使用

PMRF 提供了预训练的盲人脸图像修复模型。您可以直接通过 Hugging Face 自动下载权重并进行推理，无需手动下载检查点。

3.1 快速推理

准备好低质量 (LQ) 的人脸图像文件夹后，运行以下命令：

python inference.py \
--ckpt_path ohayonguy/PMRF_blind_face_image_restoration \
--ckpt_path_is_huggingface \
--lq_data_path /path/to/lq/images \
--output_dir /path/to/results/dir \
--batch_size 64 \
--num_flow_steps 25

参数说明：

• --lq_data_path: 输入的低质量图像目录路径。
• --output_dir: 修复结果保存目录。
• --num_flow_steps: 流步数（论文中的超参数 $K$），默认 25，可根据需求调整以平衡速度与质量。
• 注意: 该预训练模型针对正方形且已对齐的人脸图像进行了优化。若要处理包含多张人脸或未对齐的通用照片，建议使用官方的 Hugging Face Demo。

3.2 使用本地模型检查点

如果您已手动下载了模型文件（例如从 Google Drive），可使用以下命令：

python inference.py \
--ckpt_path ./checkpoints/blind_face_restoration_pmrf.ckpt \
--lq_data_path /path/to/lq/images \
--output_dir /path/to/results/dir \
--batch_size 64 \
--num_flow_steps 25