结构化移动网页的生成着色

结构化移动网页生成彩色化的正式实施，WACV 2023。

ArXiv |数据集 |预训练模型

设置

开发环境

Ubuntu 22.04、Python 3.10.9、诗歌 1.2.2
CUDA 11.6、cuDNN 8.7.0
PyTorch 1.12.1、PyTorch Lightning 1.8.6、深度图库 0.9.1
（用于截图）Google Chrome 108.0.5359.124、ChromeDriver 108.0.5359.71
（用于计算对比度违规）Lighthouse 9.6.8

安装

git clone https://github.com/CyberAgentAILab/webcolor.git
poetry install

请注意，我们无法保证或支持在其他环境（例如 Windows）中的操作。如果您希望为其他 CUDA 版本安装 PyTorch 或 DGL，请编辑 pyproject.toml 中的 URL。您可以在此处找到在 Ubuntu 上安装 Chrome、ChromeDriver 和 Lighthouse 的命令。

数据准备

./data/download.sh cache

有关数据集的详细信息，请参阅此文档。

着色演示

MODEL_NAME=CVAE  # {CVAE,NAR,AR,Stats}
BASE_URL=https://storage.googleapis.com/ailab-public/webcolor/checkpoints
poetry run python demo.py --model $MODEL_NAME --ckpt_path ${BASE_URL} / ${MODEL_NAME} .ckpt --upsampler_path ${BASE_URL} /Upsampler.ckpt --target random --out_path output/screenshot.png --num_save 3 --save_gt

上面的命令使用预先训练的模型执行自动着色，并生成如下所示的屏幕截图。

CVAE #1	CVAE #2	CVAE#3	真实的

训练

MODEL_NAME=CVAE  # {CVAE,NAR,AR,Stats,Upsampler}
poetry run python -m webcolor.main fit --model $MODEL_NAME --trainer.accelerator gpu --trainer.devices 1

模型超参数可以使用--model.help $MODEL_NAME列出。

评估

MODEL_NAME=CVAE  # {CVAE,NAR,AR,Stats,Upsampler}
CKPT_PATH=https://storage.googleapis.com/ailab-public/webcolor/checkpoints/ ${MODEL_NAME} .ckpt  # Evaluate the pre-trained model
# CKPT_PATH=lightning_logs/version_0/checkpoints/best.ckpt  # Evaluate your own trained model
poetry run python -m webcolor.main test --model $MODEL_NAME --ckpt_path $CKPT_PATH --trainer.default_root_dir /tmp --trainer.accelerator gpu --trainer.devices 1

以下命令计算 Pixel-FCD 和对比度违规，需要很长时间才能完成（在我们的环境中有 24 个工作人员，大约需要四个小时）。

MODEL_NAME=CVAE  # {CVAE,NAR,AR,Stats}
CKPT_PATH=https://storage.googleapis.com/ailab-public/webcolor/checkpoints/ ${MODEL_NAME} .ckpt
# CKPT_PATH=lightning_logs/version_0/checkpoints/best.ckpt
UPSAMPLER_PATH=https://storage.googleapis.com/ailab-public/webcolor/checkpoints/Upsampler.ckpt
# UPSAMPLER_PATH=lightning_logs/version_1/checkpoints/best.ckpt
poetry run python eval.py --num_workers 4 --model $MODEL_NAME --ckpt_path $CKPT_PATH --upsampler_path $UPSAMPLER_PATH

有关预训练模型的详细信息，请参阅此文档。

引文

 @inproceedings { Kikuchi2023 ,
    title = { Generative Colorization of Structured Mobile Web Pages } ,
    author = { Kotaro Kikuchi and Naoto Inoue and Mayu Otani and Edgar Simo-Serra and Kota Yamaguchi } ,
    booktitle = { IEEE/CVF Winter Conference on Applications of Computer Vision (WACV) } ,
    year = { 2023 } ,
    pages = { 3639-3648 } ,
    doi = { 10.1109/WACV56688.2023.00364 }
}