xypb / colorvae

EECS442 21WN Final Project: Color-VAE

[Report] [Demo Video] [Slides]

Environment setup

Please install survae first

pip install git+https://github.com/didriknielsen/survae_flows.git
# the install corresponding environment
conda create -n colorvae --file requirements.txt

Downlaod Dataset

1. COCO train2017 & val2017 dataset Please extract to ./coco folder.

2. Tiny imagenet provided on Kaggle

   kaggle datasets download -d akash2sharma/tiny-imagenet

   Please delete the duplicated folder in the zip file before training

Download Pre-trained Model

[Google Drive]

Demo

python predict.py --resume models/dil256-vae_model.pt --img_path par37351-teaser-story-big.jpg

Train

python main.py --img_size 256  --dataset COCO --lr 0.001 --exp_name dil256 --batch_size 32

Prediction

1. For image output

python predict.py --img_size 256 --resume models/dil256-vae_model.pt --img_path coco/val2017 --sample_num 8 --separate

2. For image output with VAE hint

python predict.py --img_size 256 --resume models/dil256-vae_model.pt --img_path coco/val2017 --sample_num 8 --separate --vae_hint

3. For PNSR score

python predict.py --img_size 256 --resume models/dil256-vae_model.pt --img_path coco/val2017 --sample_num 8 --psnr

Performance

On COCO val2017 dataset

PSNR

ColorVAE	No VAE	ECCV16	SIGGRAPH17	w/o semantic pre-train	w/ VAE hint
26.0008	24.6531	21.9863	25.6168	21.5979	27.0091

Rating

No VAE	ECCV16	SIGGRAPH17	w/ VAE hint
3.555	3.188	3.546	3.691