構造化されたモバイル Web ページの生成的なカラー化

構造化モバイル Web ページの生成的カラー化、WACV 2023 の正式実装。

アーシーブ |データセット |事前トレーニングされたモデル

設定

開発環境

Ubuntu 22.04、Python 3.10.9、Poetry 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 を編集してください。 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

上記のコマンドは、事前トレーニングされたモデルを使用して自動色付けを実行し、次のようなスクリーンショットを生成します。

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 人のワーカーで約 4 時間)。

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