UDiffText

UDiffText: 基于字符感知扩散模型的高质量任意图像文本合成统一框架

我们提出的UDiffText能够在合成或真实世界图像中合成准确和谐的文本，因此可以应用于场景文本编辑(a)、任意文本生成(b)和精确的文本到图像生成(c)等任务

UDiffText 预览

📬 新闻

2023.7.16 我们的论文被ECCV2024接收！🥳
2023.12.11 版本2.0更新（清理了无用代码🚮）
2023.12.3 构建Hugging Face演示
2023.12.1 构建GitHub项目页面
2023.11.30 版本1.0上传

🔨 安装

克隆此仓库：

git clone https://github.com/ZYM-PKU/UDiffText.git
cd UDiffText

安装所需的Python包

conda create -n udiff python=3.11
conda activate udiff
pip install torch==2.1.1 torchvision==0.16.1 --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt

mkdir ./checkpoints

checkpoints
├── AEs                    // 自动编码器
├── encoders             
    ├── LabelEncoder       // 字符级编码器
    └── ViTSTR             // STR编码器
├── predictors             // STR模型
├── pretrained             // 预训练SD
└── ***.ckpt               // UDiffText检查点

💻 训练

准备数据

LAION-OCR

在磁盘上创建数据目录**{你的数据根目录}/LAION-OCR并将数据放入其中。然后在./configs/dataset/locr.yaml中设置data_root**字段。
关于Laion-OCR数据集的下载和预处理，请参考TextDiffuser和我们的**./scripts/preprocess/laion_ocr_pre.ipynb**。

ICDAR13

在磁盘上创建数据目录**{你的数据根目录}/ICDAR13并将数据放入其中。然后在./configs/dataset/icd13.yaml中设置data_root**字段。
构建如下的树结构：

ICDAR13
├── train                  // 训练集
    ├── annos              // 注释
        ├── gt_x.txt
        ├── ...
    └── images             // 图像
        ├── img_x.jpg
        ├── ...
└── val                    // 验证集
    ├── annos              // 注释
        ├── gt_img_x.txt
        ├── ...
    └── images             // 图像
        ├── img_x.jpg
        ├── ...

TextSeg

在磁盘上创建数据目录**{你的数据根目录}/TextSeg并将数据放入其中。然后在./configs/dataset/tsg.yaml中设置data_root**字段。
构建如下的树结构：

TextSeg
├── train                  // 训练集
    ├── annotation         // 注释
        ├── x_anno.json    // 注释json文件
        ├── x_mask.png     // 字符级掩码
        ├── ...
    └── image              // 图像
        ├── x.jpg.jpg
        ├── ...
└── val                    // 验证集
    ├── annotation         // 注释
        ├── x_anno.json    // 注释json文件
        ├── x_mask.png     // 字符级掩码
        ├── ...
    └── image              // 图像
        ├── x.jpg
        ├── ...

SynthText

在磁盘上创建数据目录**{你的数据根目录}/SynthText并将数据放入其中。然后在./configs/dataset/st.yaml中设置data_root**字段。
构建如下的树结构：

SynthText
├── 1                      // 第1部分
    ├── ant+hill_1_0.jpg   // 图像
    ├── ant+hill_1_1.jpg
    ├── ...
├── 2                      // 第2部分
├── ...
└── gt.mat                 // 注释文件

训练字符级编码器

设置**./configs/pretrain.yaml**中的参数并运行：

python pretrain.py

训练UDiffText模型

下载预训练模型并将其放在**./checkpoints/pretrained/**中。加载检查点时可以忽略"Missing Key"或"Unexcepted Key"警告。

设置**./configs/train.yaml**中的参数，特别是路径：

load_ckpt_path: ./checkpoints/pretrained/512-inpainting-ema.ckpt // 预训练SD的检查点
model_cfg_path: ./configs/train/textdesign_sd_2.yaml // UDiffText模型配置
dataset_cfg_path: ./configs/dataset/locr.yaml // 使用Laion-OCR数据集

然后运行：

python train.py

📏 评估

下载我们可用的检查点并将它们放在**./checkpoints**中的相应目录中。
设置**./configs/test.yaml**中的参数，特别是路径：

load_ckpt_path: "./checkpoints/***.ckpt"  // UDiffText检查点
model_cfg_path: "./configs/test/textdesign_sd_2.yaml"  // UDiffText模型配置
dataset_cfg_path: "./configs/dataset/locr.yaml"  // LAION-OCR数据集配置

然后运行：

python test.py

🖼️ 演示

要在自己的机器上运行交互式演示，请执行以下代码：

python demo.py

或者在hugging face上尝试我们的在线演示：

🎉 致谢

数据集：我们衷心感谢TextDiffuser提供的具有字符级分割的开源大型图像-文本数据集LAION-OCR。
代码和模型：我们的项目基于Stable Diffusion XL的代码仓库构建，并利用了Stable Diffusion 2.0的预训练检查点。

🪬 引用

@misc{zhao2023udifftext,
      title={UDiffText: A Unified Framework for High-quality Text Synthesis in Arbitrary Images via Character-aware Diffusion Models}, 
      author={Yiming Zhao and Zhouhui Lian},
      year={2023},
      eprint={2312.04884},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}