Colorisation générative de pages Web mobiles structurées

Implémentation officielle de la colorisation générative des pages Web mobiles structurées, WACV 2023.

ArXiv | Ensemble de données | Modèles pré-entraînés

Installation

Environnement de développement

Ubuntu 22.04, Python 3.10.9, Poésie 1.2.2
CUDA 11.6, cuDNN 8.7.0
PyTorch 1.12.1, PyTorch Lightning 1.8.6, bibliothèque de graphiques profonds 0.9.1
(Pour prendre des captures d'écran) Google Chrome 108.0.5359.124, ChromeDriver 108.0.5359.71
(Pour calculer les violations de contraste) Lighthouse 9.6.8

Installation

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

Notez que nous ne pouvons pas garantir ou prendre en charge le fonctionnement dans d'autres environnements, tels que Windows. Si vous souhaitez installer PyTorch ou DGL pour d'autres versions de CUDA, veuillez modifier les URL dans pyproject.toml. Vous pouvez trouver les commandes pour installer Chrome, ChromeDriver et Lighthouse sur Ubuntu ici.

Préparation des données

./data/download.sh cache

Pour plus de détails sur l’ensemble de données, veuillez consulter ce document.

Démo de colorisation

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

La commande ci-dessus effectue une colorisation automatique à l'aide de modèles pré-entraînés et produit des captures d'écran comme celles-ci.

CVAE #1	CVAE #2	CVAE #3	Réel

Entraînement

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

Les hyperparamètres du modèle peuvent être répertoriés avec --model.help $MODEL_NAME .

Évaluation

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

La commande suivante calcule les violations Pixel-FCD et de contraste et prend beaucoup de temps (environ quatre heures avec 24 travailleurs dans notre environnement).

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

Pour plus de détails sur les modèles pré-entraînés, veuillez consulter ce document.

Citation

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