Awesome
Learning to Structure an Image with Few Colors [Website] [arXiv]
@inproceedings{hou2020learning,
title={Learning to Structure an Image with Few Colors},
author={Hou, Yunzhong and Zheng, Liang and Gould, Stephen},
booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
pages={10116--10125},
year={2020}
}
Overview
We release the PyTorch code for ColorCNN, a newly introduced architecture in our paper Learning to Structure an Image with Few Colors.
Content
Dependencies
This code uses the following libraries
- python 3.7+
- pytorch 1.4+ & torchvision
- numpy
- matplotlib
- pillow
- opencv-python
Data Preparation
By default, all datasets are in ~/Data/
. We use CIFAR10, CIFAR100, STL10, and tiny-imagenet-200 in this project.
The first three datasets can be automatically downloaded.
Tiny-imagenet-200 can be downloaded from this link.
Once downloaded, please extract the zip files under ~/Data/tiny200/
.
Then, run python color_distillation/utils/tiny_imagenet_val_reformat.py
to reformat the validation set. (thank @tjmoon0104 for his code).
Your ~/Data/
folder should look like this
Data
├── cifar10/
│ └── ...
├── cifar100/
│ └── ...
├── stl10/
│ └── ...
└── tiny200/
├── train/
│ └── ...
├── val/
│ ├── n01443537/
│ └── ...
└── ...
Code
One can find classifier training & evaluation for traditional color quantization methods in grid_downsample.py
.
For ColorCNN training & evaluation, please find it in color_cnn_downsample.py
.
Training Classifiers
In order to train classifiers, please specify '--train'
in the arguments.
python grid_downsample.py -d cifar10 -a alexnet --train
One can run the shell script bash train_classifiers.sh
to train AlexNet on all four datasets.
Training & Evaluating ColorCNN
Based on the original image pre-trained classifiers, we then train ColorCNN under specific color space sizes.
python color_cnn_downsample.py -d cifar10 -a alexnet --num_colors 2
Please run the shell script bash train_test_colorcnn.sh
to train and evaluate ColorCNN with AlexNet on all four datasets, under a 1-bit color space.
Evaluating Traditional Methods
Based on pre-trained classifiers, one can directly evaluate the performance of tradition color quantization methods.
python python grid_downsample.py -d cifar10 -a alexnet --num_colors 2 --sample_type mcut --dither
Please run the shell script bash test_mcut_dither.sh
to evaluate MedianCut+Dithering with AlexNet on all four datasets, under a 1-bit color space.