Генеративная раскраска структурированных мобильных веб-страниц

Официальная реализация генеративной раскраски структурированных мобильных веб-страниц, WACV 2023.

АрXив | Набор данных | Предварительно обученные модели

Настраивать

Среда разработки

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. Если вы хотите установить PyTorch или DGL для других версий CUDA, отредактируйте URL-адреса в pyproject.toml. Здесь вы можете найти команды для установки Chrome, ChromeDriver и Lighthouse в Ubuntu.

Подготовка данных

./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

Приведенная выше команда выполняет автоматическую раскраску с использованием предварительно обученных моделей и создает снимки экрана, подобные приведенным ниже.

КВАЭ №1	КВАЭ № 2	КВАЭ №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 }
}