forked from bnsreenu/python_for_microscopists
-
Notifications
You must be signed in to change notification settings - Fork 0
/
092-autoencoder_colorize_transfer_learning_VGG16_V0.1.py
145 lines (120 loc) · 4.96 KB
/
092-autoencoder_colorize_transfer_learning_VGG16_V0.1.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
#!/usr/bin/env python
__author__ = "Sreenivas Bhattiprolu"
__license__ = "Feel free to copy, I appreciate if you acknowledge Python for Microscopists"
# https://youtu.be/bIaT2X5Hd5k
"""
Image colorization using Autoencoders. Transfer learning using VGG.
"""
from keras.layers import Conv2D, UpSampling2D, Input
from keras.models import Sequential, Model
from keras.preprocessing.image import ImageDataGenerator, img_to_array, load_img
from skimage.color import rgb2lab, lab2rgb, gray2rgb
from skimage.transform import resize
from skimage.io import imsave
import numpy as np
import tensorflow as tf
import keras
import os
"""
Because we are going to replace the encoder part with VGG16,
we don’t need it as a classifier, we need it as a feature extractor so,
the last dense layers isn’t needed we have to pop them up.
here, we iterate on each layer except the last dense layers so,
we add 19 layer to our model. the dimension of last layer volume is “7x7x512”.
we will be using that latent space volume as a feature vector to be input to the decoder.
and the decoder is going to learn the mapping from the latent space vector to ab channels.
we want the layers of VGG16 with its original weights without changing them,
so that we set the trainable parameter in each layer to false because we don’t want to train them again.
"""
from keras.applications.vgg16 import VGG16
vggmodel = VGG16()
newmodel = Sequential()
#num = 0
for i, layer in enumerate(vggmodel.layers):
if i<19: #Only up to 19th layer to include feature extraction only
newmodel.add(layer)
newmodel.summary()
for layer in newmodel.layers:
layer.trainable=False #We don't want to train these layers again, so False.
"""
VGG16 is expecting an image of 3 dimension with size 224x224 as an input,
in preprocessing we have to scale all images to 224 instead of 256
"""
path = 'images/colorization/'
#Normalize images - divide by 255
train_datagen = ImageDataGenerator(rescale=1. / 255)
train = train_datagen.flow_from_directory(path, target_size=(224, 224), batch_size=32, class_mode=None)
#Convert from RGB to Lab
"""
by iterating on each image, we convert the RGB to Lab.
Think of LAB image as a grey image in L channel and all color info stored in A and B channels.
The input to the network will be the L channel, so we assign L channel to X vector.
And assign A and B to Y.
"""
X =[]
Y =[]
for img in train[0]:
try:
lab = rgb2lab(img)
X.append(lab[:,:,0])
Y.append(lab[:,:,1:] / 128) #A and B values range from -127 to 128,
#so we divide the values by 128 to restrict values to between -1 and 1.
except:
print('error')
X = np.array(X)
Y = np.array(Y)
X = X.reshape(X.shape+(1,)) #dimensions to be the same for X and Y
print(X.shape)
print(Y.shape)
#now we have one channel of L in each layer but, VGG16 is expecting 3 dimension,
#so we repeated the L channel two times to get 3 dimensions of the same L channel
vggfeatures = []
for i, sample in enumerate(X):
sample = gray2rgb(sample)
sample = sample.reshape((1,224,224,3))
prediction = newmodel.predict(sample)
prediction = prediction.reshape((7,7,512))
vggfeatures.append(prediction)
vggfeatures = np.array(vggfeatures)
print(vggfeatures.shape)
#Decoder
model = Sequential()
model.add(Conv2D(256, (3,3), activation='relu', padding='same', input_shape=(7,7,512)))
model.add(Conv2D(128, (3,3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(64, (3,3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(32, (3,3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(16, (3,3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))
model.add(UpSampling2D((2, 2)))
model.summary()
model.compile(optimizer='Adam', loss='mse' , metrics=['accuracy'])
model.fit(vggfeatures, Y, verbose=1, epochs=10, batch_size=128)
model.save('colorize_autoencoder_VGG16.model')
############################################
#Predicting using saved model.
model = tf.keras.models.load_model('colorize_autoencoder_VGG16_10000.model',
custom_objects=None,
compile=True)
testpath = 'images/colorization2/test_images/'
files = os.listdir(testpath)
for idx, file in enumerate(files):
test = img_to_array(load_img(testpath+file))
test = resize(test, (224,224), anti_aliasing=True)
test*= 1.0/255
lab = rgb2lab(test)
l = lab[:,:,0]
L = gray2rgb(l)
L = L.reshape((1,224,224,3))
#print(L.shape)
vggpred = newmodel.predict(L)
ab = model.predict(vggpred)
#print(ab.shape)
ab = ab*128
cur = np.zeros((224, 224, 3))
cur[:,:,0] = l
cur[:,:,1:] = ab
imsave('images/colorization2/vgg_result/result'+str(idx)+".jpg", lab2rgb(cur))