PaintsChainer server on CPU

cgi_exe.py

#!/usr/bin/env python

"""
from https://github.com/taizan/PaintsChainer/blob/master/cgi-bin/paint_x2_unet/cgi_exe.py


This code allows you to colorize on CPU (without CUDA).


$ python server.py

(another session)
$ open http://localhost:8000/static/


BUG: During colorize, RGB/BGR swap is occured... you need to change color code as follows:

in cgi_exe.py:
	change cv2.COLOR_YUV2BGR to cv2.COLOR_YUV2RGB

in img2imgDataset.py:
	change cv2.COLOR_BGR2YUV to cv2.COLOR_RGB2YUV

"""

import numpy as np
import chainer
import chainer.functions as F
import chainer.links as L
import six
import os
import cv2

from PIL import Image

from chainer import cuda, optimizers, serializers, Variable
from chainer import training
from chainer.training import extensions

from img2imgDataset import ImageAndRefDataset

import unet
class Paintor:
    def __init__(self, gpu = 0):
 
        print("start")
        self.root = "./static/images/"
        self.batchsize = 1
        self.outdir = self.root+"out/"
        self.outdir_min = self.root+"out_min/"
        self.gpu = gpu

        print("load model")
        cuda.get_device(self.gpu).use()
        self.cnn_128 = unet.UNET()
        self.cnn = unet.UNET()
        # self.cnn_128.to_gpu()
        # self.cnn.to_gpu()
        serializers.load_npz("./cgi-bin/paint_x2_unet/models/unet_128_standard", self.cnn_128)
        serializers.load_npz("./cgi-bin/paint_x2_unet/models/unet_512_standard", self.cnn)


    def save_as_img( self, array , name ):
        array = array.transpose(1,2,0)
        array = np.clip(array,0,255)
        img = np.uint8(array)
        img = cv2.cvtColor( img , cv2.COLOR_YUV2BGR )
        cv2.imwrite( name , img )
        

    def colorize( self, id_str, blur=0, s_size=128):
        cuda.get_device(self.gpu).use()

        dataset = ImageAndRefDataset([id_str+".png"],self.root+"line/",self.root+"ref/" ) 
        test_in_s, test_in = dataset.get_example(0,minimize=True)
        # 1st fixed to 128*128
        x = np.zeros((1, 4, test_in_s.shape[1], test_in_s.shape[2])).astype('f')
        input_bat = np.zeros((1, 4, test_in.shape[1], test_in.shape[2] )).astype('f')
        print(input_bat.shape)

        line ,line2 =  dataset.get_example(0,minimize=True)
        x[0,:] = line
        input_bat[0,0,:] = line2

        # x = cuda.to_gpu(x)
        x = x
        y = self.cnn_128.calc(Variable(x), test=True )

        # output = y.data.get()
        output = y.data

        self.save_as_img( output[0], self.outdir_min + id_str +"_"+ str(0) + ".png" )

        # 2nd
        for ch in range(3):
            input_bat[0,1+ch,:] = cv2.resize(output[0,ch,:], (test_in.shape[2], test_in.shape[1]), interpolation = cv2.INTER_CUBIC)

        # x = cuda.to_gpu(input_bat)
        x = input_bat
        y = self.cnn.calc(Variable(x), test=True )

        # output = y.data.get()
        output = y.data

        self.save_as_img( output[0], self.outdir + id_str +"_"+ str(0) + ".jpg" )


if __name__ == '__main__':
	paintor = Paintor( gpu=None )
	paintor.colorize('1')