def colorize_flask(self, id_str, step='C', blur=0, s_size=128, colorize_format="jpg"):
if self.gpu >= 0:
cuda.get_device(self.gpu).use()
_ = {'S': "ref/", 'L': "out_min/", 'C': "ref/"}
dataset = ImageAndRefDataset(
[id_str + ".png"], self.root + "line/", self.root + _[step])
_ = {'S': True, 'L': False, 'C': True}
sample = dataset.get_example(0, minimize=_[step], blur=blur, s_size=s_size)
_ = {'S': 0, 'L': 1, 'C': 0}[step]
# shape is 1,4,height,width
sample_container = np.zeros(
(1, 4, sample[_].shape[1], sample[_].shape[2]), dtype='f')
sample_container[0, :] = sample[_]
if self.gpu >= 0:
sample_container = cuda.to_gpu(sample_container)
cnn = {'S': self.cnn_128, 'L': self.cnn_512, 'C': self.cnn_128}
with chainer.no_backprop_mode():
with chainer.using_config('train', False):
image_conv2d_layer = cnn[step].calc(Variable(sample_container))
del sample_container
""" convert 128 to 512 """
if step == 'C':
input_bat = np.zeros((1, 4, sample[1].shape[1], sample[1].shape[2]), dtype='f')
input_bat[0, 0, :] = sample[1]
output = cuda.to_cpu(image_conv2d_layer.data[0])
del image_conv2d_layer # release memory
for channel in range(3):
input_bat[0, 1 + channel, :] = cv2.resize(
output[channel, :],
(sample[1].shape[2], sample[1].shape[1]),
interpolation=cv2.INTER_CUBIC)
if self.gpu >= 0:
link = cuda.to_gpu(input_bat, None)
else:
link = input_bat
with chainer.no_backprop_mode():
with chainer.using_config('train', False):
image_conv2d_layer = self.cnn_512.calc(Variable(link))
del link # release memory
image_out_path = {
'S': self.outdir_min + id_str + ".png",
'L': self.outdir + id_str + ".jpg",
'C': self.outdir + id_str + "_0." + colorize_format}
self.save_as_img(image_conv2d_layer.data[0], image_out_path[step])
del image_conv2d_layer
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)
y = self.cnn_128.calc(Variable(x), test=True)
output = y.data.get()
self.save_as_img(output[0],
self.outdir_min + id_str + "_" + str(0) + ".png")
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)
y = self.cnn.calc(Variable(x), test=True)
output = y.data.get()
self.save_as_img(output[0],
self.outdir + id_str + "_" + str(0) + ".jpg")
def colorize_l(self, id_str):
cuda.get_device(self.gpu).use()
dataset = ImageAndRefDataset([id_str + ".png"], self.root + "line/",
self.root + "out_min/")
test_in, test_in_ = dataset.get_example(0, minimize=False)
x = np.zeros((1, 4, test_in.shape[1], test_in.shape[2])).astype('f')
x[0, :] = test_in
x = cuda.to_gpu(x)
y = self.cnn.calc(Variable(x), test=True)
output = y.data.get()
self.save_as_img(output[0], self.outdir + id_str + ".jpg")
def colorize_l(self, id_str):
if self.gpu >= 0:
cuda.get_device(self.gpu).use()
dataset = ImageAndRefDataset([id_str + ".png"], self.root + "line/",
self.root + "out_min/")
test_in, test_in_ = dataset.get_example(0, minimize=False)
x = np.zeros((1, 4, test_in.shape[1], test_in.shape[2]), dtype='f')
x[0, :] = test_in
if self.gpu >= 0:
x = cuda.to_gpu(x)
y = self.cnn.calc(Variable(x, volatile='on'), test=True)
self.save_as_img(y.data[0], self.outdir + id_str + ".jpg")
def colorize_s(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,
blur=blur,
s_size=s_size)
x = np.zeros(
(1, 4, test_in_s.shape[1], test_in_s.shape[2])).astype('f')
x[0, :] = test_in_s
x = cuda.to_gpu(x)
y = self.cnn_128.calc(Variable(x), test=True)
output = y.data.get()
self.save_as_img(output[0], self.outdir_min + id_str + ".png")
def colorize(self, id_str, blur=0, s_size=128):
if self.gpu >= 0:
cuda.get_device(self.gpu).use()
dataset = ImageAndRefDataset([id_str + ".png"], self.root + "line/",
self.root + "ref/")
line, line2 = dataset.get_example(0, minimize=True)
# 1st fixed to 128*128
x = np.zeros((1, 4, line.shape[1], line.shape[2]), dtype='f')
input_bat = np.zeros((1, 4, line2.shape[1], line2.shape[2]), dtype='f')
print(input_bat.shape)
x[0, :] = line
input_bat[0, 0, :] = line2
if self.gpu >= 0:
x = cuda.to_gpu(x, cuda.Stream.null)
y = self.cnn_128.calc(Variable(x, volatile='on'), test=True)
del x # release memory
output = cuda.to_cpu(y.data[0])
del y # release memory
# self.save_as_img(output, self.outdir_min + id_str + "_0.png")
for ch in range(3):
input_bat[0,
1 + ch, :] = cv2.resize(output[ch, :],
(line2.shape[2], line2.shape[1]),
interpolation=cv2.INTER_CUBIC)
if self.gpu >= 0:
x = cuda.to_gpu(input_bat, cuda.Stream.null)
else:
x = input_bat
y = self.cnn.calc(Variable(x, volatile='on'), test=True)
del x # release memory
self.save_as_img(y.data[0], self.outdir + id_str + "_0.jpg")
def colorize(self, id_str, step='C', blur=0, s_size=128,colorize_format="jpg"):
if self.gpu >= 0:
cuda.get_device(self.gpu).use()
_ = {'S': "ref/", 'L': "out_min/", 'C': "ref/"}
dataset = ImageAndRefDataset(
[id_str + ".png"], self.root + "line/", self.root + _[step])
_ = {'S': True, 'L': False, 'C': True}
sample = dataset.get_example(0, minimize=_[step], blur=blur, s_size=s_size)
_ = {'S': 0, 'L': 1, 'C': 0}[step]
sample_container = np.zeros(
(1, 4, sample[_].shape[1], sample[_].shape[2]), dtype='f')
sample_container[0, :] = sample[_]
if self.gpu >= 0:
sample_container = cuda.to_gpu(sample_container)
cnn = {'S': self.cnn_128, 'L': self.cnn_512, 'C': self.cnn_128}
image_conv2d_layer = cnn[step].calc(Variable(sample_container, volatile='on'), test=True)
del sample_container
if step == 'C':
input_bat = np.zeros((1, 4, sample[1].shape[1], sample[1].shape[2]), dtype='f')
print(input_bat.shape)
input_bat[0, 0, :] = sample[1]
output = cuda.to_cpu(image_conv2d_layer.data[0])
del image_conv2d_layer # release memory
for channel in range(3):
input_bat[0, 1 + channel, :] = cv2.resize(
output[channel, :],
(sample[1].shape[2], sample[1].shape[1]),
interpolation=cv2.INTER_CUBIC)
if self.gpu >= 0:
link = cuda.to_gpu(input_bat, None)
else:
link = input_bat
image_conv2d_layer = self.cnn_512.calc(Variable(link, volatile='on'), test=True)
del link # release memory
image_out_path = {
'S': self.outdir_min + id_str + ".png",
'L': self.outdir + id_str + ".jpg",
'C': self.outdir + id_str + "_0." + colorize_format}
self.save_as_img(image_conv2d_layer.data[0], image_out_path[step])
del image_conv2d_layer
### custom
#img_out = Image.open(self.outdir + id_str +"_"+ str(0) + ".jpg")
img_out = cv2.imread(self.outdir + id_str +"_"+ str(0) + ".jpg")
# line image is original resolution
img_line = Image.open(self.root+"line/" + id_str + ".png");
img_line = img_line.convert("RGB")
img_out = cv2.resize(img_out, img_line.size, interpolation = cv2.INTER_CUBIC)
#img_out.save(self.root+"tmp_out.png")
cv2.imwrite(self.root+"tmp_out.png", img_out)
#img_out.save(self.root+ "out/"+id_str + "tmp_out.png")
cv2.imwrite(self.root+ "out/"+id_str + "tmp_out.png", img_out)
del img_out
# median out mult
img_tmp_out = cv2.imread(self.root+ "out/"+id_str + "tmp_out.png")
img_med = cv2.medianBlur(img_tmp_out, ksize=5)
cv2.imwrite(self.root+"median_out.png", img_med)
cv2.imwrite(self.root+ "out/"+id_str + "median_out.png", img_med)
img_out = Image.open(self.root+ "out/"+id_str + "median_out.png")
img_mult = ImageChops.multiply(img_out, img_line)
img_mult.save(self.root+"med_mult.png")
img_mult.save(self.root+ "out/"+id_str + "med_mult.png")
# multiply
#img_out = Image.open(self.root+"tmp_out.png")
#img_mult = ImageChops.multiply(img_out, img_line)
#img_mult.save(self.root+"mult.png")
del img_mult
# output only color
img_line = ImageOps.invert(img_line)
#img_line = img_line.resize(img_out.size)
#img_sub = ImageChops.subtract(img_out, img_line)
img_add = ImageChops.add(img_out, img_line)
img_add.save(self.root+"color.png")
img_add.save(self.root+ "out/"+id_str + "color.png")
del img_out, img_line, img_add
# median to color
img_color = cv2.imread(self.root+ "out/"+id_str + "color.png")
img_med = cv2.medianBlur(img_color, ksize=5)
cv2.imwrite(self.root+"median_color.png", img_med)
cv2.imwrite(self.root+ "out/"+id_str + "median_color.png", img_med)
# median color mult
#img_line = Image.open(self.root+"line/" + id_str + ".png");
#img_line = img_line.convert("RGB")
#img_out = Image.open(self.root+"median_color.png")
#img_mult = ImageChops.multiply(img_out, img_line)
#img_mult.save(self.root+"med_mult_color.png")
# median color soft
img_med = cv2.imread(self.root+ "out/"+id_str + "median_color.png")
img_soft = cv2.GaussianBlur(img_med, (15,15), 1)
cv2.imwrite(self.root+"median_color_soft.png", img_soft)
cv2.imwrite(self.root+ "out/"+id_str + "median_color_soft.png", img_soft)
def colorize(self, id_str, step='C', blur=0, s_size=128, colorize_format="jpg"):
if self.gpu >= 0:
cuda.get_device(self.gpu).use()
dataset = ImageAndRefDataset(
[id_str + ".png"], self.root + "line/", self.root + "ref/",
colormode=self._colormode, normalized=self._norm)
# s_size is the size of final result, could be modified
sample = dataset.get_example(0, minimize=True, blur=blur, s_size=128)
if self._colormode == 'YUV':
sample_container = np.zeros(
(1, 4, sample[0].shape[1], sample[0].shape[2]), dtype='f')
elif self._colormode == 'LAB':
sample_container = np.zeros(
(1, 3, sample[0].shape[1], sample[0].shape[2]), dtype='f')
sample_container[0, :] = sample[0]
if self.gpu >= 0:
sample_container = cuda.to_gpu(sample_container)
image_conv2d_layer = self.cnn_128.calc(Variable(sample_container, volatile='on'), test=True)
del sample_container
# if step == 'C':
# input_bat = np.zeros((1, 4, sample[1].shape[1], sample[1].shape[2]), dtype='f')
# print(input_bat.shape)
# input_bat[0, 0, :] = sample[1]
# output = cuda.to_cpu(image_conv2d_layer.data[0])
# del image_conv2d_layer # release memory
# for channel in range(3):
# input_bat[0, 1 + channel, :] = cv2.resize(
# output[channel, :],
# (sample[1].shape[2], sample[1].shape[1]),
# interpolation=cv2.INTER_CUBIC)
# if self.gpu >= 0:
# link = cuda.to_gpu(input_bat, None)
# else:
# link = input_bat
# image_conv2d_layer = self.cnn_512.calc(Variable(link, volatile='on'), test=True)
# del link # release memory
image_out_path = self.outdir + id_str + "_0." + colorize_format
if self._colormode == 'YUV':
convdata = image_conv2d_layer.data[0]
self.save_as_img(convdata, image_out_path)
elif self._colormode == 'LAB':
if self._norm:
convdata = image_conv2d_layer.data[0] + 128
else:
convdata = image_conv2d_layer.data[0]
ab = convdata.transpose(1, 2, 0)
ab = ab.clip(0, 255).astype(np.uint8)
if self.gpu >= 0:
ab = cuda.to_cpu(ab)
L = sample[1].transpose(1, 2, 0)
lab = np.zeros((L.shape[0], L.shape[1], 3), dtype='f')
lab[:,:,0] = L[:,:,0]
lab[:,:,1:] = cv2.resize(ab, (L.shape[1], L.shape[0]), interpolation=cv2.INTER_AREA)
lab = lab.astype(np.uint8)
res = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
cv2.imwrite(image_out_path, res)
del image_conv2d_layer