Bridging the Domain Gap towards Generalization in Automatic Colorization

This is pytorch impelementation of "Bridging the Domain Gap towards Generalization in Automatic Colorization", Hyejin Lee, Deahee Kim, DaEun Lee, Jinkyu Kim, Jaekoo Lee, ECCV 2022 [paper] (update soon)

Overview

Main Architecture

Output

Dataset

We use PACS, Office-Home Dataset. Dataset Structure is provided below.
We split examples from training domains in the ratio 8:2 (train:validation)
Please prepare following this structure.

If you want to split dataset, you can refer to our code (utils.py - split_domain_data function)

PACS Dataset : https://domaingeneralization.github.io/

Office-Home Dataset : https://www.hemanthdv.org/officeHomeDataset.html

|--Dataset root
|   |-- Domain1 (ex: Photo)
|   |   |-- Class1 (ex: Dog)
|   |   |   |-- img1.jpg
|   |   |   |-- ...
|   |   |-- ...
|   |-- Domain2
|   |   |--Class1
|   |   |--...
|   |-- ...
|   |-- Train
|   |   |-- Domain1
|   |   |   |-- Class1
|   |   |   |   |-- img1.jpg
|   |   |   |   |--...
|   |   |   |--...
|   |   |-- Domain2
|   |   |   |-- Class1
|   |   |-- ...
|   |-- Valid
|   |   |-- Domain1
|   |   |   |-- Class1
|   |   |   |   |-- img1.jpg
|   |   |   |   |-- ...
|   |   |   |-- ...
|   |   |-- Domain2
|   |   |   |-- Class1
|   |   |-- ...

Requirements & Environment

we used one RTX 3090 GPU for model training. and we also tested V100 GPU and A6000 GPU.

pip install -r requirements.txt

In particular, ensure that the version of Pytorch 1.7+, python 3.8+, cuda 11.0+

Train

Run the following command to train the DG colorization Model

python train.py --config CONFIG --save_dir SAVE_DIR --w-adv [coefficient of adv loss] --gpu [gpu number]

Test

Run the following comman to test the model

python test.py --h_model_path [head encoder model.pth] --t_model_path [tail encoder model.pth] --d_model_path [decoder model.pth] \
               --source [source domain name] --domain [target domain name] \
               --data_dir [testset root] --save_dir  [path to save] \
               --gpu [gpu number]

Dataset Structure is provided below. Please input dataset root path in 'data_dir' argument.

|--Dataset dataset root
|   |-- Domain1 (ex: Photo)
|   |   |-- Class1 (ex: Dog)
|   |   |   |-- img1.jpg
|   |   |   |-- ...
|   |   |-- ...
|   |-- Domain2
|   |   |--Class1
|   |   |--...
|   |-- ...

Citation

update soon