Skip to content

Pix2pix that allows users to design the color of generated images

Notifications You must be signed in to change notification settings

soap117/Color_pix2pix

Repository files navigation

Color_pix2pix

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.
Image

Requirement

python 3.5
tesorflow 1.0
opencv3
PIL or pillow
scipy

Getting Stated

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\\

Training the model

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.
Image
Forexample:

python train_gray2real.py
Enter the training data location: E:\\color_pix2pix\\handbags\\
Enter the name of save files: handbags

Testing the model

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

Generating your own data sets

The first is the same as in pix2pix-tensorflow. Then use the data_pre.py to generate the training set for color_pix2pix.

About

Pix2pix that allows users to design the color of generated images

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages