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
• Data Preparation
• Code
  • Training Classifiers
  • Training & Evaluating ColorCNN
  • Evaluating Traditional Methods

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.