通过去噪扩散模型合成人像

新闻

2023.02 可以通过Google Colab使用一个演示版本：

:rocket: 在Colab上演示

生成结果

你可以直接从Google Drive下载我们的测试结果：(1) PIDM.zip (2) PIDM_vs_Others.zip

文件 PIDM_vs_Others.zip 比较了我们的方法与几种最先进的方法，例如ADGAN [14]、PISE [24]、GFLA [20]、DPTN [25]、CASD [29]、NTED [19]。每行包含 target_pose、source_image、ground_truth、ADGAN、PISE、GFLA、DPTN、CASD、NTED 和 PIDM（我们的）。

数据集

从In-shop Clothes Retrieval Benchmark 下载DeepFashion数据集的 img_highres.zip。
解压 img_highres.zip。你需要向数据集维护者询问密码。然后将获得的文件夹重命名为 img 并放在 ./dataset/deepfashion 目录下。
我们按照 GFLA 的方法划分训练/测试集。有几个具有显著遮挡的图片从训练集中移除。下载下列文件以获取训练/测试对及通过 Openpose 提取的关键点 pose.zip：
- 从 Google Drive 下载训练/测试对，包括 train_pairs.txt、test_pairs.txt、train.lst、test.lst。把这些文件放在 ./dataset/deepfashion 目录下。
- 从 Google Drive 下载通过Openpose提取的关键点 pose.rar。解压并将获得的文件夹放在 ./dataset/deepfashion 目录下。

运行以下代码将图片保存到lmdb数据集：

python data/prepare_data.py \
--root ./dataset/deepfashion \
--out ./dataset/deepfashion

自定义数据集

任何自定义数据集的文件夹结构应如下：

dataset/
- <dataset_name>/
- - img/
- - pose/
- - train_pairs.txt
- - test_pairs.txt

你的所有图片基本上会放在 img 文件夹中。你可以使用不同的子文件夹来存储你的图片，或将所有图片放在 img 文件夹内。对应的姿势存储在 pose 文件夹中（如果你使用openpose，则为txt文件。在我们的项目中，我们使用18点关键点估计）。train_pairs.txt 和 test_pairs.txt 包含所有可能对的路径，以逗号分隔 <src_path1>,<tgt_path1>。

之后，运行以下命令处理数据：

python data/prepare_data.py \
--root ./dataset/<dataset_name> \
--out ./dataset/<dataset_name>
--sizes ((256,256),)

这将创建一个lmdb数据集 ./dataset/<dataset_name>/256-256/

Conda 安装

# 1. 创建一个conda虚拟环境。
conda create -n PIDM python=3.7
conda activate PIDM
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia

# 2. 克隆仓库并安装依赖
git clone https://github.com/ankanbhunia/PIDM
pip install -r requirements.txt

方法

训练

此代码支持多GPU训练。在DeepFashion数据集上，使用8个A100 GPU和批量大小为8的情况下，完整训练需要5天时间。该模型在300个epoch后训练完成；然而，在200个epoch后它便可以生成高质量的可用样本。我们也尝试使用V100 GPU进行训练，我们的代码训练时间相似。

python -m torch.distributed.launch --nproc_per_node=8 --master_port 48949 train.py \
--dataset_path "./dataset/deepfashion" --batch_size 8 --exp_name "pidm_deepfashion"

推理

从这里下载预训练模型并将其放在 checkpoints 文件夹中。对于姿势控制，使用 obj.predict_pose，如下代码片段：

from predict import Predictor
obj = Predictor()

obj.predict_pose(image=<PATH_OF_SOURCE_IMAGE>, sample_algorithm='ddim', num_poses=4, nsteps=50)

对于外观控制，使用 obj.predict_appearance

from predict import Predictor
obj = Predictor()

src = <PATH_OF_SOURCE_IMAGE>
ref_img = <PATH_OF_REF_IMAGE>
ref_mask = <PATH_OF_REF_MASK>
ref_pose = <PATH_OF_REF_POSE>

obj.predict_appearance(image=src, ref_img = ref_img, ref_mask = ref_mask, ref_pose = ref_pose, sample_algorithm = 'ddim',  nsteps = 50)

输出将保存为文件名 output.png。

引用

如果你在你的研究中使用了该结果和代码，请引用我们的论文：

@article{bhunia2022pidm,
  title={Person Image Synthesis via Denoising Diffusion Model},
  author={Bhunia, Ankan Kumar and Khan, Salman and Cholakkal, Hisham and Anwer, Rao Muhammad and Laaksonen, Jorma and Shah, Mubarak and Khan, Fahad Shahbaz},
  journal={CVPR},
  year={2023}
}

Ankan Kumar Bhunia, Salman Khan, Hisham Cholakkal, Rao Anwer, Jorma Laaksonen, Mubarak Shah & Fahad Khan