基于去噪扩散模型的人物图像合成

ArXiv | 项目主页 | 演示 | YouTube 视频

最新消息

• 2023年2月 通过 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]。每一行分别包含目标姿态、源图像、真实标签、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/
• • <数据集名称>/
• • • 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/<数据集名称> \
--out ./dataset/<数据集名称>
--sizes ((256,256),)

这将创建一个 lmdb 格式的数据集 ./dataset/<数据集名称>/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=<源图像路径>, sample_algorithm='ddim', num_poses=4, nsteps=50)

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

from predict import Predictor
obj = Predictor()

src = <源图像路径>
ref_img = <参考图像路径>
ref_mask = <参考掩码路径>
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

PIDM 快速上手指南

PIDM (Person Image Synthesis via Denoising Diffusion Model) 是一个基于去噪扩散模型的人物图像合成工具，支持姿态控制（Pose Control）和外观控制（Appearance Control）。

环境准备

• 操作系统: Linux (推荐 Ubuntu)
• Python 版本: 3.7
• GPU 要求: 支持 CUDA 的 NVIDIA 显卡（训练推荐多卡 A100/V100，推理单卡即可）
• 前置依赖: Conda 包管理器

安装步骤

1. 创建虚拟环境并安装 PyTorch

建议使用国内镜像源加速下载。

# 创建名为 PIDM 的 conda 环境
conda create -n PIDM python=3.7
conda activate PIDM

# 安装 PyTorch (使用清华镜像源加速)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
# 或者使用官方命令 (如果网络通畅)
# conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia

2. 克隆项目并安装依赖

git clone https://github.com/ankanbhunia/PIDM
cd PIDM
pip install -r requirements.txt

基本使用

1. 准备预训练模型

从 Google Drive 下载预训练模型文件，并将其放入项目根目录下的 checkpoints 文件夹中。

2. 运行推理 (Inference)

PIDM 提供两种主要的生成模式：姿态控制和外观控制。请确保已准备好相应的输入图片（源图像、参考图像、姿态文件等）。

模式一：姿态控制 (Pose Control)

根据源图像生成指定数量不同姿态的新图像。

from predict import Predictor

obj = Predictor()

# image: 源图像路径
# num_poses: 生成不同姿态的数量
# nsteps: 采样步数
obj.predict_pose(image=<PATH_OF_SOURCE_IMAGE>, sample_algorithm='ddim', num_poses=4, nsteps=50)

模式二：外观控制 (Appearance Control)

将源图像的人物外观迁移到参考图像的姿态上。

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。

---

注：如需使用自定义数据集或进行训练，请参考项目原文档配置 DeepFashion 数据集结构及运行 train.py。