【亲测免费】 PMRF项目使用与启动教程

PMRF项目使用与启动教程

1. 项目介绍

PMRF（Posterior-Mean Rectified Flow）是一个新颖的图片复原算法，该算法可以证明近似最优估计器，在完美的感知质量约束下最小化均方误差（MSE）。本项目是ICLR 2025会议论文的官方实现，旨在实现最小MSE的图片复原，同时保证图片的逼真度。

2. 项目快速启动

环境搭建

首先，你需要创建一个conda环境并安装所需的依赖：

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

运行示例

为了快速使用模型，我们提供了Hugging Face checkpoint。以下是一个快速运行的示例：

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

如果你想要使用本地模型checkpoint，可以按以下方式运行：

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 demo。

3. 应用案例和最佳实践

案例一：盲脸图片复原

使用PMRF模型进行盲脸图片复原时，可以参考以下步骤：

1. 准备低质量的图片作为输入。
2. 使用inference.py脚本进行图片复原。
3. 可以使用提供的Hugging Face checkpoint或本地checkpoint。

最佳实践

• 确保输入图片的尺寸符合模型训练时的尺寸要求。
• 根据具体应用场景调整num_flow_steps参数以获得最佳效果。

4. 典型生态项目

PMRF作为图像复原领域的一个先进算法，可以与其他图像处理项目配合使用，例如：

• 使用PMRF进行图像预处理，然后输入到其他图像识别或生成模型中。
• 结合PMRF和其他图像质量评估工具，如lpips和piq，进行综合图像质量评估。

以上就是PMRF项目的基本介绍和快速启动指南。希望对你有所帮助！