zeruniverse / neural-colorization Goto Github PK

View Code? Open in Web Editor NEW

156.0 5.0 40.0 40 KB

(Pytorch Implementation) GAN for image colorization

License: Other

Python 100.00%

neural-colorization's Introduction

neural-colorization

GAN for image colorization based on Johnson's network structure.

Setup

Install the following Python libraries:

numpy
scipy
Pytorch
scikit-image
Pillow
opencv-python

Colorize images

#Download pre-trained model
wget -O model.pth "https://github.com/zeruniverse/neural-colorization/releases/download/1.1/G.pth"

#Colorize an image with CPU
python colorize.py -m model.pth -i input.jpg -o output.jpg --gpu -1

# If you want to colorize all images in a folder with GPU
python colorize.py -m model.pth -i input -o output --gpu 0

Train your own model

Note: Training is only supported with GPU (CUDA).

Prepare dataset

Download some datasets and unzip them into a same folder (saying train_raw_dataset). If the images are not in .jpg format, you should convert them all in .jpgs.
run python build_dataset_directory.py -i train_raw_dataset -o train (you can skip this if all your images are directly under the train_raw_dataset, in which case, just rename the folder as train)
run python resize_all_imgs.py -d train to resize all training images into 256*256 (you can skip this if your images are already in 256*256)

Optional preparation

It's highly recommended to train from my pretrained models. You can get both generator model and discriminator model from the GitHub Release:

wget "https://github.com/zeruniverse/neural-colorization/releases/download/1.1/G.pth"
wget "https://github.com/zeruniverse/neural-colorization/releases/download/1.1/D.pth"

It's also recommended to have a test image (the script will generate a colorization for the test image you give at every checkpoint so you can see how the model works during training).

Training

The required arguments are training image directory (e.g. train) and path to save checkpoints (e.g. checkpoints)

python train.py -d train -c chekpoints

To add initial weights and test images:

python train.py -d train -c chekpoints --d_init D.pth --g_init G.pth -t test.jpg

More options are available and you can run python train.py --help to print all options.

For torch equivalent (no GAN), you can set option -p 1e9 (set a very large weight for pixel loss).

Reference

Perceptual Losses for Real-Time Style Transfer and Super-Resolution

License

GNU GPL 3.0 for personal or research use. COMMERCIAL USE PROHIBITED.

Model weights are released under CC BY 4.0

neural-colorization's People

Contributors

Stargazers

Watchers

neural-colorization's Issues

Dataset

Can you please tell me where you train your network?

About Cuda deserialize object

There is a video card in my computer, but when I use the following command, I am prompted that it cannot be mapped。
python colorize.py -m G.pth -i input.jpg -o output.jpg --gpu -1

python colorize.py -m G.pth -i input.jpg -o output.jpg --gpu 0

Error information:

RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map
your storages to an existing device.

what's version of your numpy

i meet the error AttributeError: 'numpy.ufunc' object has no attribute 'module'
what's version of your numpy or caused by other reasons,thanks

Epoch: 3: [D loss: 50.000000] [G total loss: 16.745750] [G GAN Loss: 0.000000]
Epoch: 3: [D loss: 50.000000] [G total loss: 17.254082] [G GAN Loss: 0.000000]
Epoch: 3: [D loss: 50.000000] [G total loss: 16.371378] [G GAN Loss: 0.000000]
Epoch: 3: [D loss: 50.000000] [G total loss: 14.244134] [G GAN Loss: 0.000000]
Epoch: 3: [D loss: 50.000000] [G total loss: 17.944151] [G GAN Loss: 0.000000]
Epoch: 3: [D loss: 50.000000] [G total loss: 20.372141] [G GAN Loss: 0.000000]
Epoch: 3: [D loss: 50.000000] [G total loss: 17.622280] [G GAN Loss: 0.000000]
Epoch: 3: [D loss: 50.000000] [G total loss: 14.337464] [G GAN Loss: 0.000000]
How can I solve it? Thank you very much!

Dataset

Can you tell me on which dataset was the model trained?

Using AdaptiveAvgPool2d(2)

What if we use AdaptiveAvgPool2d(2) inorder to avoid the size mismatch error of resnet18 and use the pre-trained model as a checkpoint? It will be like this,
D.fc = nn.Sequential(nn.Linear(2048, 1), nn.Sigmoid())
D.avgpool = nn.AdaptiveAvgPool2d(2)

AdaptiveAvgPool2d

Self Training Throw error, mis matching in shape

Hi, I am using this repo to train a image colorization model for a project. And I tried to use train.py to do some training using pretrained model D and G. The errors come out as:
-> g_loss_gan = adversarial_loss(D(gen_imgs), valid) (Pdb) RuntimeError: size mismatch, m1: [20 x 512], m2: [2048 x 1] at /opt/conda/conda-bld/pytorch_1535491974311/work/aten/src/THC/generic/THCTensor MathBlas.cu:249
Where
(Pdb) type(D) <class 'torchvision.models.resnet.ResNet'>
and
(Pdb) gen_imgs.shape torch.Size([20, 3, 256, 256])
The images are shaped into 256*256
Is anything incorrect here?

Porting this network to deeplearn.js (or TensorFlow.js)

Hello! I really liked Your project, but I have no experience with Torch, Python or Linux at all. Can you help me and convert your project to DeepLearn.JS or TensorFlow.JS (these libraries were developed by Google and can be found on GitHub. They are completely written in JavaScript.). In fact, I am interested in the weight of your neural network. I work in the Windows operating system and study neural networks in C#.

Getting the weights of a given neural network

Hello! I would love to run Your project, but I can not do it. Could you extract the weights of this neural network and save them to a binary file(as a sequence of numbers like Single) so that I can run the network?

Asking Info

On how many images was it trained? And also how many epochs did you use to generate the pre-trained model?

"python train.py -d train -c chekpoints --d_init D.pth --g_init G.pth -t test.jpg" error

thanks for your great work. but i am running your code "python train.py -d train -c chekpoints --d_init D.pth --g_init G.pth -t test.jpg" for tanning will pop below errors. would you please fix it? thanks

File "train.py", line 140, in
D.load_state_dict(model_dict)
File "/usr/local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 839, in load_state_dict
self.class.name, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for ResNet:
size mismatch for fc.0.weight: copying a param with shape torch.Size([1, 2048]) from checkpoint, the shape in current model is torch.Size([1, 512]).

Implement pytorch version and GAN

Use GAN to remove gray artifacts

Translate to pytorch WIP:
pytorch.zip