Home

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. system overview

Content

Dependencies

This code uses the following libraries

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.