PITI：预训练就是图像到图像翻译所需的一切

官方 PyTorch 实现

预训练就是图像到图像翻译所需的一切
王腾飞、张婷、张博、欧阳浩、陈栋、陈启峰、温芳
2022年

论文 | 项目网站 | 视频 | 在线演示

简介

我们提出了一种简单且通用的框架，将预训练的力量应用于各种图像到图像的翻译任务。如果您感兴趣，可以尝试我们的在线演示。

由我们的方法合成的多样化样本。

搭建环境

安装

git clone https://github.com/PITI-Synthesis/PITI.git
cd PITI

环境

conda env create -f environment.yml

快速入门

预训练模型

请下载我们的 Base 模型和 Upsample 模型的预训练权重，并将其放置在 ./ckpt 目录下。

模型	任务	数据集
Base-64x64	掩码到图像	在 COCO 数据集上训练。
Upsample-64-256	掩码到图像	在 COCO 数据集上训练。
Base-64x64	草图到图像	在 COCO 数据集上训练。
Upsample-64-256	草图到图像	在 COCO 数据集上训练。

如果无法访问这些链接，您也可以从这里获取我们的预训练模型。

准备图像

我们在 ./test_imgs 中放置了一些示例图像，您可以快速试用它们。

COCO

对于 COCO 数据集，您可以从COCO 官网下载图像和标注文件。

对于掩码到图像的合成，我们使用 RGB 格式的语义图作为输入。要获得这样的语义图，请运行 ./preprocess/preprocess_mask.py（原始掩码和处理后的掩码示例位于 preprocess/example 中）。请注意，我们不需要像之前的工作那样即时生成的掩码。

对于草图到图像的合成，我们使用通过 HED 提取的草图作为输入。要获得这样的草图，请运行 ./preprocess/preprocess_sketch.py。

推理

交互式推理

运行以下脚本，它将创建一个基于 Gradio 的交互式 GUI。您可以上传输入掩码或草图并生成图像。

pip install gradio
python inference.py

批量推理

根据以下说明修改 sample.sh，然后运行：

bash sample.sh

参数	描述
--model_path	基础模型检查点的路径。
--sr_model_path	上采样模型检查点的路径。
--val_data_dir	包含图像路径的文本文件的路径。
--num_samples	您希望采样的图像数量。
--sample_c	无分类器指导的强度。
--mode	输入类型。

训练

准备工作

1. 下载并预处理数据集。对于 COCO 数据集，从COCO 官网下载图像和标注文件。运行 ./preprocess/preprocess_mask.py 或 ./preprocess/preprocess_sketch.py
2. 使用 python preprocess/download.py 下载预训练模型。

开始训练

以掩码到图像合成为例：（草图到图像相同）

微调基础模型

修改 mask_finetune_base.sh 并运行：

bash mask_finetune_base.sh

微调上采样模型

修改 mask_finetune_upsample.sh 并运行：

bash mask_finetune_upsample.sh

引用

如果您发现这项工作对您的研究有帮助，请引用：

@article{wang2022pretraining,
 title = {Pretraining is All You Need for Image-to-Image Translation},
  author = {Wang, Tengfei and Zhang, Ting and Zhang, Bo and Ouyang, Hao and Chen, Dong and Chen, Qifeng and Wen, Fang},
  journal={arXiv:2205.12952},
  year = {2022},
}

致谢

感谢 GLIDE 分享他们的代码。

PITI 快速上手指南

PITI (Pretraining is All You Need for Image-to-Image Translation) 是一个简单且通用的图像到图像翻译框架，支持从语义掩码（Mask）或草图（Sketch）生成高质量图像。

环境准备

• 操作系统: Linux (推荐 Ubuntu 18.04+)
• 硬件要求: NVIDIA GPU (建议显存 8GB 以上)，已安装 CUDA 驱动
• 前置依赖:
  • Python 3.8+
  • Conda (推荐使用 Miniconda 或 Anaconda 管理环境)
  • Git

安装步骤

1. 克隆项目代码

git clone https://github.com/PITI-Synthesis/PITI.git
cd PITI

2. 创建虚拟环境

使用项目提供的配置文件一键安装依赖：

conda env create -f environment.yml
conda activate PITI

提示：如果 environment.yml 下载依赖较慢，可手动创建环境后，使用国内镜像源（如清华源）安装核心包： pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 下载预训练模型

本项目需要下载 Base (基础生成) 和 Upsample (超分放大) 两类模型。请将下载的 .pt 文件放入 ./ckpt 目录。

任务类型	模型类型	分辨率	下载链接 (需自行替换为可用源或原链接)
Mask-to-Image	Base	64x64	下载链接
Mask-to-Image	Upsample	64->256	下载链接
Sketch-to-Image	Base	64x64	下载链接
Sketch-to-Image	Upsample	64->256	下载链接

注：若上述链接无法访问，可尝试在项目的备用文件夹中查找。

基本使用

方式一：交互式界面 (推荐新手)

最简单的使用方式是启动 Gradio 网页界面，上传掩码或草图即可实时生成图像。

1. 安装 gradio：

pip install gradio

2. 运行推理脚本：

python inference.py

运行后终端会显示本地访问地址（通常为 http://127.0.0.1:7860），在浏览器打开即可使用。

方式二：命令行批量推理

如需处理本地图片文件夹，可修改并运行 sample.sh 脚本。

1. 编辑 sample.sh，确保以下参数路径正确：

   • --model_path: Base 模型路径
   • --sr_model_path: Upsample 模型路径
   • --val_data_dir: 包含输入图片路径的 txt 文件
   • --mode: 输入类型 (mask 或 sketch)

2. 执行脚本：

bash sample.sh

数据预处理说明

如果你使用自己的数据集（如 COCO），需要先进行预处理：

• Mask-to-Image: 运行 ./preprocess/preprocess_mask.py 将语义图转换为 RGB 格式输入。
• Sketch-to-Image: 运行 ./preprocess/preprocess_sketch.py 提取 HED 草图特征。

项目已在 ./test_imgs 目录下提供示例图片，可直接用于测试。