This porject is based on pix2pix and the code is based on pix2pix-tensorflow. The original pix2pix doesn't have the ability to choose the color. So I made a little reform
in the original pix2pix network. In color_pix2pix, the network is breaked into 2 parts. The first part generates the gray images according to the edges pictures. The second part is aimed at coloring the gray pictures according to a rough color distribution.
In addition, the second part of net is trained in directly minimizing the difference between generated images and groud turth images because I found this approach is better than adversarial training in coloring the images.
python 3.5
tesorflow 1.0
opencv3
PIL or pillow
scipy
Download the data sets from the https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/.
Then extract the data to your local disk.
Run the /data/data_pre.py to generate the training data.
Forexample:
cd "the location of data"
python data_pre.py
Enter the original data location: E:\\edges2handbags\\train\\
Where you want to save the new data: E:\\color_pix2pix\\handbags\\
The models will be saved in ./logs.
Run the train_edage2gray.py to train the first part of model.
Forexample:
python train_edage2gray.py
Enter the training data location: E:\\color_pix2pix\\handbags\\
Enter the name of save files: handbags
Run the train_gray2real.py to train the second part of model.
In second part, we resize the ground truth into a very small size like 6*6.
Then resize it into original size to obtain a rough color distribution.
After this operation, the color distribution only contains the information of color.
Forexample:
python train_gray2real.py
Enter the training data location: E:\\color_pix2pix\\handbags\\
Enter the name of save files: handbags
Run the eval.py to test model.
The pre-trained model can be donwloaded from the logs file.
Download the logs.zip files and extract them.
2 models of handbags and shoes are uploaded into the logs file.
Forexample:
python eval.py
Enter the name of save files: handbags
Enter the name of line picture: handbag.png
Enter the name of color picture: handbag_c.png
The first is the same as in pix2pix-tensorflow. Then use the data_pre.py to generate the training set for color_pix2pix.