Colorização generativa de páginas web estruturadas para dispositivos móveis

Implementação oficial da Colorização Generativa de Páginas Web Estruturadas para Dispositivos Móveis, WACV 2023.

ArXiv | Conjunto de dados | Modelos pré-treinados

Configurar

Ambiente de desenvolvimento

Ubuntu 22.04, Python 3.10.9, Poesia 1.2.2
CUDA 11.6, cuDNN 8.7.0
PyTorch 1.12.1, PyTorch Lightning 1.8.6, Biblioteca Deep Graph 0.9.1
(Para fazer capturas de tela) Google Chrome 108.0.5359.124, ChromeDriver 108.0.5359.71
(Para calcular violações de contraste) Lighthouse 9.6.8

Instalação

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

Observe que não podemos garantir ou oferecer suporte à operação em outros ambientes, como o Windows. Se você deseja instalar PyTorch ou DGL para outras versões CUDA, edite os URLs em pyproject.toml. Você pode encontrar os comandos para instalar o Chrome, ChromeDriver e Lighthouse no Ubuntu aqui.

Preparação de dados

./data/download.sh cache

Para obter detalhes sobre o conjunto de dados, consulte este documento.

Demonstração de colorização

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

O comando acima realiza colorização automática usando modelos pré-treinados e produz capturas de tela como as seguintes.

AVC #1	AVC #2	AVC #3	Real

Treinamento

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

Os hiperparâmetros do modelo podem ser listados com --model.help $MODEL_NAME .

Avaliação

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

O comando a seguir calcula Pixel-FCD e violações de contraste e leva muito tempo para ser concluído (cerca de quatro horas com 24 trabalhadores em nosso ambiente).

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

Para obter detalhes sobre os modelos pré-treinados, consulte este documento.

Citação

 @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 }
}