-
Notifications
You must be signed in to change notification settings - Fork 51
/
main.py
113 lines (92 loc) · 4.16 KB
/
main.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
from keras.models import Sequential
from keras.layers import Conv2D, Input, BatchNormalization
# from keras.layers.advanced_activations import LeakyReLU
from keras.callbacks import ModelCheckpoint
from keras.optimizers import SGD, Adam
import prepare_data as pd
import numpy
import math
def psnr(target, ref):
# assume RGB image
target_data = numpy.array(target, dtype=float)
ref_data = numpy.array(ref, dtype=float)
diff = ref_data - target_data
diff = diff.flatten('C')
rmse = math.sqrt(numpy.mean(diff ** 2.))
return 20 * math.log10(255. / rmse)
def model():
# lrelu = LeakyReLU(alpha=0.1)
SRCNN = Sequential()
SRCNN.add(Conv2D(nb_filter=128, nb_row=9, nb_col=9, init='glorot_uniform',
activation='relu', border_mode='valid', bias=True, input_shape=(32, 32, 1)))
SRCNN.add(Conv2D(nb_filter=64, nb_row=3, nb_col=3, init='glorot_uniform',
activation='relu', border_mode='same', bias=True))
# SRCNN.add(BatchNormalization())
SRCNN.add(Conv2D(nb_filter=1, nb_row=5, nb_col=5, init='glorot_uniform',
activation='linear', border_mode='valid', bias=True))
adam = Adam(lr=0.0003)
SRCNN.compile(optimizer=adam, loss='mean_squared_error', metrics=['mean_squared_error'])
return SRCNN
def predict_model():
# lrelu = LeakyReLU(alpha=0.1)
SRCNN = Sequential()
SRCNN.add(Conv2D(nb_filter=128, nb_row=9, nb_col=9, init='glorot_uniform',
activation='relu', border_mode='valid', bias=True, input_shape=(None, None, 1)))
SRCNN.add(Conv2D(nb_filter=64, nb_row=3, nb_col=3, init='glorot_uniform',
activation='relu', border_mode='same', bias=True))
# SRCNN.add(BatchNormalization())
SRCNN.add(Conv2D(nb_filter=1, nb_row=5, nb_col=5, init='glorot_uniform',
activation='linear', border_mode='valid', bias=True))
adam = Adam(lr=0.0003)
SRCNN.compile(optimizer=adam, loss='mean_squared_error', metrics=['mean_squared_error'])
return SRCNN
def train():
srcnn_model = model()
print(srcnn_model.summary())
data, label = pd.read_training_data("./train.h5")
val_data, val_label = pd.read_training_data("./test.h5")
checkpoint = ModelCheckpoint("SRCNN_check.h5", monitor='val_loss', verbose=1, save_best_only=True,
save_weights_only=False, mode='min')
callbacks_list = [checkpoint]
srcnn_model.fit(data, label, batch_size=128, validation_data=(val_data, val_label),
callbacks=callbacks_list, shuffle=True, nb_epoch=200, verbose=0)
# srcnn_model.load_weights("m_model_adam.h5")
def predict():
srcnn_model = predict_model()
srcnn_model.load_weights("3051crop_weight_200.h5")
IMG_NAME = "/home/mark/Engineer/SR/data/Set14/flowers.bmp"
INPUT_NAME = "input2.jpg"
OUTPUT_NAME = "pre2.jpg"
import cv2
img = cv2.imread(IMG_NAME, cv2.IMREAD_COLOR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
shape = img.shape
Y_img = cv2.resize(img[:, :, 0], (shape[1] / 2, shape[0] / 2), cv2.INTER_CUBIC)
Y_img = cv2.resize(Y_img, (shape[1], shape[0]), cv2.INTER_CUBIC)
img[:, :, 0] = Y_img
img = cv2.cvtColor(img, cv2.COLOR_YCrCb2BGR)
cv2.imwrite(INPUT_NAME, img)
Y = numpy.zeros((1, img.shape[0], img.shape[1], 1), dtype=float)
Y[0, :, :, 0] = Y_img.astype(float) / 255.
pre = srcnn_model.predict(Y, batch_size=1) * 255.
pre[pre[:] > 255] = 255
pre[pre[:] < 0] = 0
pre = pre.astype(numpy.uint8)
img = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
img[6: -6, 6: -6, 0] = pre[0, :, :, 0]
img = cv2.cvtColor(img, cv2.COLOR_YCrCb2BGR)
cv2.imwrite(OUTPUT_NAME, img)
# psnr calculation:
im1 = cv2.imread(IMG_NAME, cv2.IMREAD_COLOR)
im1 = cv2.cvtColor(im1, cv2.COLOR_BGR2YCrCb)[6: -6, 6: -6, 0]
im2 = cv2.imread(INPUT_NAME, cv2.IMREAD_COLOR)
im2 = cv2.cvtColor(im2, cv2.COLOR_BGR2YCrCb)[6: -6, 6: -6, 0]
im3 = cv2.imread(OUTPUT_NAME, cv2.IMREAD_COLOR)
im3 = cv2.cvtColor(im3, cv2.COLOR_BGR2YCrCb)[6: -6, 6: -6, 0]
print "bicubic:"
print cv2.PSNR(im1, im2)
print "SRCNN:"
print cv2.PSNR(im1, im3)
if __name__ == "__main__":
train()
predict()