def train(self, config):
if config.is_train:
input_setup(self.sess, config)
else:
nx, ny = input_setup(self.sess, config) # 合并图像块数
if config.is_train:
data_path = os.path.join("./", config.checkpoint_dir, "train.h5")
else:
data_path = os.path.join("./", config.checkpoint_dir, "test.h5")
train_data, train_label = read_data(data_path)
self.train_op = tf.train.GradientDescentOptimizer(
config.learning_rate).minimize(self.loss)
tf.global_variables_initializer().run()
counter = 0 # 输出判断数
start_time = time.time()
# 加载训练数据
if self.load(config.checkpoint_dir):
print("[*] Load SUCCESS")
else:
print("[!] Load Failed")
if config.is_train:
print("Train....")
batch_index = len(train_data) // config.batch_size
for ep in range(config.epoch):
for idx in range(batch_index):
batch_images = train_data[idx *
config.batch_size:(idx + 1) *
config.batch_size]
batch_labels = train_label[idx *
config.batch_size:(idx + 1) *
config.batch_size]
_, err = self.sess.run([self.train_op, self.loss], {
self.images: batch_images,
self.labels: batch_labels
})
counter += 1
if counter % 10 == 0:
print(
"Epoch: %2d,step: %2d,time: %4.4f,loss: %.8f" %
((ep + 1), counter, time.time() - start_time, err))
if counter % 500 == 0:
self.save(config.checkpoint_dir, counter)
else:
print("Test...")
result = self.pred.eval({
self.images: train_data,
self.labels: train_label
})
result = merge(result, [nx, ny])
result = result.squeeze() # squeese():把 result 的 ? 维度删除
image_path = os.path.join(os.getcwd(), config.sample_dir,
"text_image.png")
imsave(image_path, result)
def test(self, test_reader):
count = 0
res = []
for corrupt, original in test_reader.read_mat(1):
x = test_reader.im2col(corrupt[0], self.patch_sz).T
y = test_reader.im2col(original[0], self.patch_sz).T
feed = {self.input_corrupt: x, self.input_original: y}
cleared_patch = self.sess.run([self.h_], feed_dict = feed)
cleared = np.clip(test_reader.reconstruct(cleared_patch[0].T, self.patch_sz, original[0].shape), 0, 1)
res.append(util.calcPSNR(cleared, original[0]))
print "%d, PSNR change %g -> %g"%(count, util.calcPSNR(corrupt[0], original[0]), res[-1])
util.imsave(cleared, str(count))
count += 1
print "avg %g, std %g"%(np.mean(res), np.std(res))