def do_for_evaluate(self, file_path, print_console=False):
true_image = util.set_image_alignment(
util.load_image(file_path, print_console=False), self.scale)
if true_image.shape[2] == 3 and self.channels == 1:
# for color images
input_y_image = loader.build_input_image(true_image,
channels=self.channels,
scale=self.scale,
alignment=self.scale,
convert_ycbcr=True)
true_y_image = util.convert_rgb_to_y(true_image)
input_bicubic_y_image = util.resize_image_by_pil(
input_y_image,
self.scale,
resampling_method=self.resampling_method)
output_y_image = self.do(input_y_image, input_bicubic_y_image)
mse = util.compute_mse(true_y_image,
output_y_image,
border_size=self.psnr_calc_border_size)
elif true_image.shape[2] == 1 and self.channels == 1:
# for monochrome images
input_image = loader.build_input_image(true_image,
channels=self.channels,
scale=self.scale,
alignment=self.scale)
input_bicubic_y_image = util.resize_image_by_pil(
input_image,
self.scale,
resampling_method=self.resampling_method)
output_image = self.do(input_image, input_bicubic_y_image)
mse = util.compute_mse(true_image,
output_image,
border_size=self.psnr_calc_border_size)
else:
mse = 0
if print_console:
print("MSE:%f, PSNR:%f" % (mse, util.get_psnr(mse)))
return mse
def do_for_file(self, file_path, output_folder="output"):
org_image = util.load_image(file_path)
filename, extension = os.path.splitext(os.path.basename(file_path))
output_folder += "/" + self.name + "/"
util.save_image(output_folder + filename + extension, org_image)
scaled_image = util.resize_image_by_pil(
org_image, self.scale, resampling_method=self.resampling_method)
util.save_image(output_folder + filename + "_bicubic" + extension,
scaled_image)
if len(org_image.shape
) >= 3 and org_image.shape[2] == 3 and self.channels == 1:
input_y_image = util.convert_rgb_to_y(org_image)
scaled_image = util.resize_image_by_pil(
input_y_image,
self.scale,
resampling_method=self.resampling_method)
util.save_image(
output_folder + filename + "_bicubic_y" + extension,
scaled_image)
output_y_image = self.do(input_y_image)
util.save_image(output_folder + filename + "_result_y" + extension,
output_y_image)
scaled_ycbcr_image = util.convert_rgb_to_ycbcr(
util.resize_image_by_pil(org_image, self.scale,
self.resampling_method))
image = util.convert_y_and_cbcr_to_rgb(
output_y_image, scaled_ycbcr_image[:, :, 1:3])
else:
scaled_image = util.resize_image_by_pil(
org_image,
self.scale,
resampling_method=self.resampling_method)
util.save_image(
output_folder + filename + "_bicubic_y" + extension,
scaled_image)
image = self.do(org_image)
util.save_image(output_folder + filename + "_result" + extension,
image)
def build_image_set(file_path,
channels=1,
scale=1,
convert_ycbcr=True,
resampling_method="bicubic",
print_console=True):
true_image = util.set_image_alignment(
util.load_image(file_path, print_console=print_console), scale)
if channels == 1 and true_image.shape[2] == 3 and convert_ycbcr:
true_image = util.convert_rgb_to_y(true_image)
input_image = util.resize_image_by_pil(true_image,
1.0 / scale,
resampling_method=resampling_method)
input_interpolated_image = util.resize_image_by_pil(
input_image, scale, resampling_method=resampling_method)
return input_image, input_interpolated_image, true_image
def build_input_batch(self):
for i in range(self.batch_num):
image = None
while image is None:
image_no = random.randrange(self.train.file_counts)
image = loader.load_random_patch(
self.train.filenames[image_no],
self.batch_image_size * self.scale,
self.batch_image_size * self.scale, self.jpeg_mode)
if random.randrange(2) == 0:
image = np.fliplr(image)
self.batch_input[i] = util.resize_image_by_pil(
image, 1 / self.scale)
self.batch_input_bicubic[i] = util.resize_image_by_pil(
self.batch_input[i], self.scale)
self.batch_true[i] = image
self.steps_in_epoch += 1
def load_batch_image(self, max_value):
""" index won't be used. """
image = None
while image is None:
image = self.load_random_patch(
self.filenames[self.get_next_image_no()])
if random.randrange(2) == 0:
image = np.fliplr(image)
input_image = util.resize_image_by_pil(image, 1 / self.scale)
input_bicubic_image = util.resize_image_by_pil(input_image, self.scale)
if max_value != 255:
scale = max_value / 255.0
input_image = np.multiply(input_image, scale)
input_bicubic_image = np.multiply(input_bicubic_image, scale)
image = np.multiply(image, scale)
return input_image, input_bicubic_image, image
def do(self, input_image, bicubic_input_image=None):
h, w = input_image.shape[:2]
ch = input_image.shape[2] if len(input_image.shape) > 2 else 1
if self.max_value != 255.0:
input_image = np.multiply(input_image, self.max_value /
255.0) # type: np.ndarray
if bicubic_input_image is not None:
bicubic_image = bicubic_input_image
elif self.bicubic_init:
bicubic_image = util.resize_image_by_pil(
input_image, self.scale, resampling_method="bicubic")
else:
bicubic_image = util.resize_image_by_pil(
input_image, self.scale, resampling_method="nearest")
quad_image = np.zeros([1, h, w,
self.output_channels]) # type: np.ndarray
loader.convert_to_multi_channel_image(quad_image[0], bicubic_image,
self.scale)
y = self.sess.run(self.y_,
feed_dict={
self.x: input_image.reshape(1, h, w, ch),
self.x2: quad_image,
self.dropout_input: 1.0
})
if self.max_value != 255.0:
quad_image = np.multiply(y[0], 255.0 / self.max_value)
else:
quad_image = y[0]
image = np.zeros(shape=[h * self.scale, w * self.scale,
1]) # type: np.ndarray
loader.convert_from_multi_channel_image(image, quad_image, self.scale)
return image
def predict_im(self, org_image):
if len(org_image.shape
) >= 3 and org_image.shape[2] == 3 and self.channels == 1:
input_y_image = util.convert_rgb_to_y(org_image)
output_y_image = self.do(input_y_image)
scaled_ycbcr_image = util.convert_rgb_to_ycbcr(
util.resize_image_by_pil(org_image, self.scale,
self.resampling_method))
image = util.convert_y_and_cbcr_to_rgb(
output_y_image, scaled_ycbcr_image[:, :, 1:3])
else:
image = self.do(org_image)
return image
def do_util(src, dest):
with Timer('util: load'):
inp = util.load_image(src)
with Timer('util: resize'):
resized = util.resize_image_by_pil(inp, 2)
with Timer('util: extract Y'):
only_y = util.convert_rgb_to_y(inp)
only_y = util.convert_rgb_to_y(resized) # simulate upscale
with Timer('util: rgb => YCbCr'):
scaled_ycbcr_image = util.convert_rgb_to_ycbcr(resized)
with Timer('util: Y + YCbCr -> rgb'):
image = util.convert_y_and_cbcr_to_rgb(only_y, scaled_ycbcr_image[:, :,
1:3])
with Timer('util: save'):
util.save_image(dest, image)
def do_for_file(self, file_path, output_folder="output"):
filename, extension = os.path.splitext(file_path)
output_folder += "/"
org_image = util.load_image(file_path)
util.save_image(output_folder + file_path, org_image)
if len(org_image.shape
) >= 3 and org_image.shape[2] == 3 and self.channels == 1:
input_y_image = util.convert_rgb_to_y(org_image,
jpeg_mode=self.jpeg_mode)
scaled_image = util.resize_image_by_pil(input_y_image, self.scale)
util.save_image(
output_folder + filename + "_bicubic_y" + extension,
scaled_image)
output_y_image = self.do(input_y_image)
util.save_image(output_folder + filename + "_result_y" + extension,
output_y_image)
scaled_ycbcr_image = util.convert_rgb_to_ycbcr(
util.resize_image_by_pil(org_image, self.scale),
jpeg_mode=self.jpeg_mode)
image = util.convert_y_and_cbcr_to_rgb(output_y_image,
scaled_ycbcr_image[:, :,
1:3],
jpeg_mode=self.jpeg_mode)
else:
scaled_image = util.resize_image_by_pil(org_image, self.scale)
util.save_image(
output_folder + filename + "_bicubic_y" + extension,
scaled_image)
image = self.do(org_image)
util.save_image(output_folder + filename + "_result" + extension,
image)
return 0
def load_input_image(self, filename, rescale=False, model=None, resampling_method="bicubic"): image = load_input_image(filename, channels=self.channels, scale=self.scale, alignment=self.alignment, jpeg_mode=self.jpeg_mode, print_console=True) if self.max_value != 255.0: image = np.multiply(image, self.max_value / 255.0) if rescale: if model is not None: rescaled_image = model.do(image) else: rescaled_image = util.resize_image_by_pil(image, self.scale, resampling_method=resampling_method) return image, rescaled_image else: return image
def do_for_evaluate(self, file_path, output_directory="output", output=True, print_console=True):
filename, extension = os.path.splitext(file_path)
output_directory += "/"
true_image = util.set_image_alignment(util.load_image(file_path), self.scale)
if true_image.shape[2] == 3 and self.channels == 1:
input_y_image = loader.build_input_image(true_image, channels=self.channels, scale=self.scale,
alignment=self.scale, convert_ycbcr=True, jpeg_mode=self.jpeg_mode)
# for color images
if output:
input_bicubic_y_image = util.resize_image_by_pil(input_y_image, self.scale)
true_ycbcr_image = util.convert_rgb_to_ycbcr(true_image, jpeg_mode=self.jpeg_mode)
output_y_image = self.do(input_y_image, input_bicubic_y_image)
mse = util.compute_mse(true_ycbcr_image[:, :, 0:1], output_y_image, border_size=self.scale)
loss_image = util.get_loss_image(true_ycbcr_image[:, :, 0:1], output_y_image, border_size=self.scale)
output_color_image = util.convert_y_and_cbcr_to_rgb(output_y_image, true_ycbcr_image[:, :, 1:3],
jpeg_mode=self.jpeg_mode)
util.save_image(output_directory + file_path, true_image)
util.save_image(output_directory + filename + "_input" + extension, input_y_image)
util.save_image(output_directory + filename + "_input_bicubic" + extension, input_bicubic_y_image)
util.save_image(output_directory + filename + "_true_y" + extension, true_ycbcr_image[:, :, 0:1])
util.save_image(output_directory + filename + "_result" + extension, output_y_image)
util.save_image(output_directory + filename + "_result_c" + extension, output_color_image)
util.save_image(output_directory + filename + "_loss" + extension, loss_image)
else:
true_y_image = util.convert_rgb_to_y(true_image, jpeg_mode=self.jpeg_mode)
output_y_image = self.do(input_y_image)
mse = util.compute_mse(true_y_image, output_y_image, border_size=self.scale)
elif true_image.shape[2] == 1 and self.channels == 1:
# for monochrome images
input_image = loader.build_input_image(true_image, channels=self.channels, scale=self.scale, alignment=self.scale)
output_image = self.do(input_image)
mse = util.compute_mse(true_image, output_image, border_size=self.scale)
if output:
util.save_image(output_directory + file_path, true_image)
util.save_image(output_directory + filename + "_result" + extension, output_image)
if print_console:
print("MSE:%f PSNR:%f" % (mse, util.get_psnr(mse)))
return mse
def upscale(self, org_image):
assert len(org_image.shape
) >= 3 and org_image.shape[2] == 3 and self.channels == 1
input_y_image = util.convert_rgb_to_y(org_image)
big_blurry_input_image = util.resize_image_by_pil(
org_image, self.scale)
big_blurry_input_ycbcr_image = util.convert_rgb_to_ycbcr(
big_blurry_input_image)
bbi_input_y_image, bbi_input_cbcr_image = util.convert_ycbcr_to_y_cbcr(
big_blurry_input_ycbcr_image)
output_y_image = self.do(input_y_image, bbi_input_y_image)
output_image = util.convert_y_and_cbcr_to_rgb(output_y_image,
bbi_input_cbcr_image)
return output_image
def make_input_image(self, file_path, true_image, scale=1, print_console=True):
if true_image is not None:
assert len(true_image.shape) == 3 and true_image.shape[2] == 3
fname, fext = os.path.splitext(file_path)
input_image = None
if fname.lower().endswith(TRUE_SUFFIX):
target = fname[:-len(TRUE_SUFFIX)] + INPUT_SUFFIX + fext
head, tail = os.path.split(target)
target = os.path.join(head, 'input', tail)
if os.path.exists(target):
input_image = util.set_image_alignment(util.load_image(target, print_console=print_console), scale)
if input_image is None and true_image is not None:
input_image = util.resize_image_by_pil(true_image, 1.0 / scale, resampling_method=self.resampling_method)
hblur_radius = random.randrange(0, self.hblur_max) / 100.
vblur_radius = random.randrange(0, self.vblur_max) / 100.
qua = random.randrange(self.jpegify[0], self.jpegify[1])
# if vblur_radius > 0:
# input_image = gaussian_filter1d(input_image, sigma=vblur_radius, axis=0)
# if hblur_radius > 0:
# input_image = gaussian_filter1d(input_image, sigma=hblur_radius, axis=1)
kx = int((hblur_radius + 0.5) * 2)
if kx % 2 == 0:
kx += 1
ky = int((vblur_radius + 0.5) * 2)
if ky % 2 == 0:
ky += 1
cv2.GaussianBlur(src=input_image, ksize=(kx, ky), sigmaX=hblur_radius, dst=input_image, sigmaY=vblur_radius)
if qua < 100:
tmpname = tempfile.mktemp(suffix='.jpg')
util.save_image(tmpname, input_image, qua, print_console=False)
input_image = util.load_image(tmpname, print_console=False)
os.unlink(tmpname)
return input_image
def do_for_file(self, file_path, output_folder="output"):
org_image = cv2.imread(file_path)
assert len(org_image.shape
) >= 3 and org_image.shape[2] == 3 and self.channels == 1
input_y_image = util.convert_rgb_to_y(org_image)
big_blurry_input_image = util.resize_image_by_pil(
org_image, self.scale)
big_blurry_input_ycbcr_image = util.convert_rgb_to_ycbcr(
big_blurry_input_image)
bbi_input_y_image, bbi_input_cbcr_image = util.convert_ycbcr_to_y_cbcr(
big_blurry_input_ycbcr_image)
output_y_image = self.do(input_y_image, bbi_input_y_image)
output_image = util.convert_y_and_cbcr_to_rgb(output_y_image,
bbi_input_cbcr_image)
target_path = os.path.basename(file_path)
cv2.imwrite(os.path.join(output_folder, target_path), output_image)
def evaluate_bicubic(self, file_path, print_console=False):
true_image = util.set_image_alignment(util.load_image(file_path, print_console=False), self.scale)
if true_image.shape[2] == 3 and self.channels == 1:
input_image = loader.build_input_image(true_image, channels=self.channels, scale=self.scale,
alignment=self.scale, convert_ycbcr=True)
true_image = util.convert_rgb_to_y(true_image)
elif true_image.shape[2] == 1 and self.channels == 1:
input_image = loader.build_input_image(true_image, channels=self.channels, scale=self.scale,
alignment=self.scale)
else:
return None, None
input_bicubic_image = util.resize_image_by_pil(input_image, self.scale, resampling_method=self.resampling_method)
psnr, ssim = util.compute_psnr_and_ssim(true_image, input_bicubic_image, border_size=self.psnr_calc_border_size)
if print_console:
print("PSNR:%f, SSIM:%f" % (psnr, ssim))
return psnr, ssim
def do(self, input_image, bicubic_input_image=None):
h, w = input_image.shape[:2]
ch = input_image.shape[2] if len(input_image.shape) > 2 else 1
if self.max_value != 255.0:
input_image = np.multiply(input_image, self.max_value / 255.0) # type: np.ndarray
if bicubic_input_image is None:
bicubic_input_image = util.resize_image_by_pil(input_image, self.scale,
resampling_method=self.resampling_method)
if self.self_ensemble > 1:
output = np.zeros([self.scale * h, self.scale * w, 1])
for i in range(self.self_ensemble):
image = util.flip(input_image, i)
bicubic_image = util.flip(bicubic_input_image, i)
y = self.sess.run(self.y_, feed_dict={self.x: image.reshape(1, image.shape[0], image.shape[1], ch),
self.x2: bicubic_image.reshape(1, self.scale * image.shape[0],
self.scale * image.shape[1],
ch),
self.dropout: 1.0, self.is_training: 0})
restored = util.flip(y[0], i, invert=True)
output += restored
output /= self.self_ensemble
else:
y = self.sess.run(self.y_, feed_dict={self.x: input_image.reshape(1, h, w, ch),
self.x2: bicubic_input_image.reshape(1, self.scale * h,
self.scale * w, ch),
self.dropout: 1.0, self.is_training: 0})
output = y[0]
if self.max_value != 255.0:
hr_image = np.multiply(output, 255.0 / self.max_value)
else:
hr_image = output
return hr_image
def load_batch_image(self, max_value):
image = None
file_path = None
while image is None:
file_path = self.filenames[self.get_next_image_no()]
if self.image_maker.patch_scale_max > 1.0:
patch_scale_value = random.randrange(100, self.image_maker.patch_scale_max * 100) / 100.
else:
patch_scale_value = 1.
image = self.load_random_patch(file_path, patch_scale_value)
input_image = self.image_maker.make_input_image(file_path, image, scale=self.scale,
print_console=False)
flip = random.randrange(0, 4)
if flip == 1 or flip == 3:
input_image = np.flipud(input_image)
image = np.flipud(image)
if flip == 2 or flip == 3:
input_image = np.fliplr(input_image)
image = np.fliplr(image)
rot90 = random.randrange(0, 2)
if rot90 == 1:
input_image = np.rot90(input_image)
image = np.rot90(image)
input_image = util.convert_rgb_to_y(input_image)
input_bicubic_image = util.resize_image_by_pil(input_image, self.scale)
if max_value != 255:
scale = max_value / 255.0
input_image = np.multiply(input_image, scale)
input_bicubic_image = np.multiply(input_bicubic_image, scale)
image = np.multiply(image, scale)
image = util.convert_rgb_to_y(image)
return input_image, input_bicubic_image, image
def do_for_file(self, file_path, output_folder="output"):
org_image = util.load_image(file_path)
filename, extension = os.path.splitext(os.path.basename(file_path))
output_folder += "/" + self.name + "/"
# util.save_image(output_folder + filename + extension, org_image)
if os.path.exists(output_folder + filename + extension):
print("File already exists in the target directory")
return
if org_image.shape[0] + org_image.shape[1] >= 1024:
print("Image is too big: ", org_image.shape)
image = org_image
elif len(org_image.shape
) >= 3 and org_image.shape[2] == 3 and self.channels == 1:
input_y_image = util.convert_rgb_to_y(org_image)
# scaled_image = util.resize_image_by_pil(input_y_image, self.scale, resampling_method=self.resampling_method)
# util.save_image(output_folder + filename + "_bicubic_y" + extension, scaled_image)
output_y_image = self.do(input_y_image)
# util.save_image(output_folder + filename + "_result_y" + extension, output_y_image)
scaled_ycbcr_image = util.convert_rgb_to_ycbcr(
util.resize_image_by_pil(org_image, self.scale,
self.resampling_method))
image = util.convert_y_and_cbcr_to_rgb(
output_y_image, scaled_ycbcr_image[:, :, 1:3])
else:
# scaled_image = util.resize_image_by_pil(org_image, self.scale, resampling_method=self.resampling_method)
# util.save_image(output_folder + filename + "_bicubic_y" + extension, scaled_image)
image = self.do(org_image)
# util.save_image(output_folder + filename + "_result" + extension, image)
file_path = output_folder + filename + extension
util.save_image(file_path, image)
print("Saved at ", file_path)
def build_input_image(image,
width=0,
height=0,
channels=1,
scale=1,
alignment=0,
convert_ycbcr=True,
jpeg_mode=False):
"""
build input image from file.
crop, adjust the image alignment for the scale factor, resize, convert color space.
"""
if width != 0 and height != 0:
if image.shape[0] != height or image.shape[1] != width:
x = (image.shape[1] - width) // 2
y = (image.shape[0] - height) // 2
image = image[y:y + height, x:x + width, :]
if image.shape[2] >= 4:
image = image[:, :, 0:3]
if alignment > 1:
image = util.set_image_alignment(image, alignment)
if scale != 1:
image = util.resize_image_by_pil(image, 1.0 / scale)
if channels == 1 and image.shape[2] == 3:
if convert_ycbcr:
image = util.convert_rgb_to_y(image, jpeg_mode=jpeg_mode)
else:
if convert_ycbcr:
image = util.convert_rgb_to_ycbcr(image, jpeg_mode=jpeg_mode)
return image
def do_for_evaluate_with_output(self,
file_path,
output_directory=None,
print_console=False):
true_image = util.set_image_alignment(
util.load_image(file_path, print_console=False), self.scale)
# Assuming the image is color
assert true_image.shape[
2] == 3 and self.channels == 1, "Only 3-channel images are supported"
input_image = loader.build_input_image(true_image,
scale=self.scale,
alignment=self.scale)
input_y_image = util.convert_rgb_to_y(input_image)
true_y_image = util.convert_rgb_to_y(true_image)
input_bicubic_y_image = util.resize_image_by_pil(
input_y_image,
self.scale,
resampling_method=self.resampling_method)
output_y_image = self.do(input_y_image, input_bicubic_y_image)
psnr, ssim = util.compute_psnr_and_ssim(
true_y_image,
output_y_image,
border_size=self.psnr_calc_border_size)
if output_directory:
true_ycbcr_image = util.convert_rgb_to_ycbcr(true_image)
_, true_cbcr = util.convert_ycbcr_to_y_cbcr(true_ycbcr_image)
output_color_image = util.convert_y_and_cbcr_to_rgb(
output_y_image, true_cbcr)
loss_image = util.get_loss_image(
true_y_image,
output_y_image,
border_size=self.psnr_calc_border_size)
filename, extension = os.path.splitext(file_path)
output_directory += "/" + self.name + "/"
util.make_dir(output_directory)
util.save_image(output_directory + file_path, true_image)
util.save_image(output_directory + filename + "_input" + extension,
input_y_image)
util.save_image(
output_directory + filename + "_input_bicubic" + extension,
input_bicubic_y_image)
util.save_image(
output_directory + filename + "_true_y" + extension,
true_ycbcr_image[:, :, 0:1])
util.save_image(
output_directory + filename + "_result" + extension,
output_y_image)
util.save_image(
output_directory + filename + "_result_c" + extension,
output_color_image)
util.save_image(output_directory + filename + "_loss" + extension,
loss_image)
if print_console:
print("[%s] PSNR:%f, SSIM:%f" % (filename, psnr, ssim))
return psnr, ssim
def log_to_tensorboard(self, test_filename, psnr, save_meta_data=True):
# todo
save_meta_data = False
org_image = util.set_image_alignment(
util.load_image(test_filename, print_console=False), self.scale)
if len(org_image.shape
) >= 3 and org_image.shape[2] == 3 and self.channels == 1:
org_image = util.convert_rgb_to_y(org_image)
input_image = util.resize_image_by_pil(
org_image,
1.0 / self.scale,
resampling_method=self.resampling_method)
feed_dict = {
self.x:
input_image.reshape([
1, input_image.shape[0], input_image.shape[1],
input_image.shape[2]
]),
self.y:
org_image.reshape([
1, org_image.shape[0], org_image.shape[1], org_image.shape[2]
]),
self.is_training:
0
}
if save_meta_data:
run_metadata = tf.RunMetadata()
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
summary_str, _ = self.sess.run([self.summary_op, self.mae],
feed_dict=feed_dict,
options=run_options,
run_metadata=run_metadata)
self.test_writer.add_run_metadata(run_metadata,
"step%d" % self.epochs_completed)
filename = self.checkpoint_dir + "/" + self.name + "_metadata.txt"
with open(filename, "w") as out:
out.write(str(run_metadata))
tf.contrib.tfprof.model_analyzer.print_model_analysis(
tf.get_default_graph(),
run_meta=run_metadata,
tfprof_options=tf.contrib.tfprof.model_analyzer.
PRINT_ALL_TIMING_MEMORY)
else:
summary_str, _ = self.sess.run([self.summary_op, self.mae],
feed_dict=feed_dict)
self.train_writer.add_summary(summary_str, self.epochs_completed)
util.log_scalar_value(self.train_writer, 'training_PSNR',
self.training_psnr_sum / self.training_step,
self.epochs_completed)
util.log_scalar_value(self.train_writer, 'LR', self.lr,
self.epochs_completed)
self.train_writer.flush()
util.log_scalar_value(self.test_writer, 'PSNR', psnr,
self.epochs_completed)
self.test_writer.flush()
def do_for_evaluate_with_output(self,
file_path,
output_directory,
print_console=False):
filename, extension = os.path.splitext(file_path)
output_directory += "/" + self.name + "/"
util.make_dir(output_directory)
true_image = util.set_image_alignment(
util.load_image(file_path, print_console=False), self.scale)
if true_image.shape[2] == 3 and self.channels == 1:
# for color images
input_y_image = loader.build_input_image(true_image,
channels=self.channels,
scale=self.scale,
alignment=self.scale,
convert_ycbcr=True)
input_bicubic_y_image = util.resize_image_by_pil(
input_y_image,
self.scale,
resampling_method=self.resampling_method)
true_ycbcr_image = util.convert_rgb_to_ycbcr(true_image)
output_y_image = self.do(input_y_image, input_bicubic_y_image)
mse = util.compute_mse(true_ycbcr_image[:, :, 0:1],
output_y_image,
border_size=self.psnr_calc_border_size)
loss_image = util.get_loss_image(
true_ycbcr_image[:, :, 0:1],
output_y_image,
border_size=self.psnr_calc_border_size)
output_color_image = util.convert_y_and_cbcr_to_rgb(
output_y_image, true_ycbcr_image[:, :, 1:3])
util.save_image(output_directory + file_path, true_image)
util.save_image(output_directory + filename + "_input" + extension,
input_y_image)
util.save_image(
output_directory + filename + "_input_bicubic" + extension,
input_bicubic_y_image)
util.save_image(
output_directory + filename + "_true_y" + extension,
true_ycbcr_image[:, :, 0:1])
util.save_image(
output_directory + filename + "_result" + extension,
output_y_image)
util.save_image(
output_directory + filename + "_result_c" + extension,
output_color_image)
util.save_image(output_directory + filename + "_loss" + extension,
loss_image)
elif true_image.shape[2] == 1 and self.channels == 1:
# for monochrome images
input_image = loader.build_input_image(true_image,
channels=self.channels,
scale=self.scale,
alignment=self.scale)
input_bicubic_y_image = util.resize_image_by_pil(
input_image,
self.scale,
resampling_method=self.resampling_method)
output_image = self.do(input_image, input_bicubic_y_image)
mse = util.compute_mse(true_image,
output_image,
border_size=self.psnr_calc_border_size)
util.save_image(output_directory + file_path, true_image)
util.save_image(
output_directory + filename + "_result" + extension,
output_image)
else:
mse = 0
if print_console:
print("[%s] MSE:%f, PSNR:%f" % (filename, mse, util.get_psnr(mse)))
return mse
def log_to_tensorboard(self, test_filename, psnr, save_meta_data=True):
if self.enable_log is False:
return
# todo
save_meta_data = False
org_image = util.set_image_alignment(
util.load_image(test_filename, print_console=False), self.scale)
if len(org_image.shape
) >= 3 and org_image.shape[2] == 3 and self.channels == 1:
org_image = util.convert_rgb_to_y(org_image)
input_image = util.resize_image_by_pil(
org_image,
1.0 / self.scale,
resampling_method=self.resampling_method)
bicubic_image = util.resize_image_by_pil(
input_image, self.scale, resampling_method=self.resampling_method)
if self.max_value != 255.0:
input_image = np.multiply(input_image, self.max_value /
255.0) # type: np.ndarray
bicubic_image = np.multiply(bicubic_image, self.max_value /
255.0) # type: np.ndarray
org_image = np.multiply(org_image,
self.max_value / 255.0) # type: np.ndarray
feed_dict = {
self.x:
input_image.reshape([
1, input_image.shape[0], input_image.shape[1],
input_image.shape[2]
]),
self.x2:
bicubic_image.reshape([
1, bicubic_image.shape[0], bicubic_image.shape[1],
bicubic_image.shape[2]
]),
self.y:
org_image.reshape([
1, org_image.shape[0], org_image.shape[1], org_image.shape[2]
]),
self.dropout:
1.0,
self.is_training:
0
}
if save_meta_data:
# profiler = tf.profiler.Profile(self.sess.graph)
run_metadata = tf.RunMetadata()
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
summary_str, _ = self.sess.run([self.summary_op, self.loss],
feed_dict=feed_dict,
options=run_options,
run_metadata=run_metadata)
self.test_writer.add_run_metadata(run_metadata,
"step%d" % self.epochs_completed)
filename = self.checkpoint_dir + "/" + self.name + "_metadata.txt"
with open(filename, "w") as out:
out.write(str(run_metadata))
# filename = self.checkpoint_dir + "/" + self.name + "_memory.txt"
# tf.profiler.write_op_log(
# tf.get_default_graph(),
# log_dir=self.checkpoint_dir,
# #op_log=op_log,
# run_meta=run_metadata)
tf.contrib.tfprof.model_analyzer.print_model_analysis(
tf.get_default_graph(),
run_meta=run_metadata,
tfprof_options=tf.contrib.tfprof.model_analyzer.
PRINT_ALL_TIMING_MEMORY)
else:
summary_str, _ = self.sess.run([self.summary_op, self.loss],
feed_dict=feed_dict)
self.train_writer.add_summary(summary_str, self.epochs_completed)
if not self.use_l1_loss:
if self.training_step != 0:
util.log_scalar_value(
self.train_writer, 'PSNR',
self.training_psnr_sum / self.training_step,
self.epochs_completed)
util.log_scalar_value(self.train_writer, 'LR', self.lr,
self.epochs_completed)
self.train_writer.flush()
util.log_scalar_value(self.test_writer, 'PSNR', psnr,
self.epochs_completed)
self.test_writer.flush()
def do_for_evaluate(self,
file_path,
output_directory=None,
print_console=False,
save_output_images=True):
if save_output_images:
filename, extension = os.path.splitext(file_path)
output_directory += "/" + self.name + "/"
util.make_dir(output_directory)
true_image = util.set_image_alignment(
util.load_image(file_path, print_console=False), self.scale)
if true_image.shape[2] == 3 and self.channels == 1:
# for color images
input_y_image = loader.build_input_image(true_image,
channels=self.channels,
scale=self.scale,
alignment=self.scale,
convert_ycbcr=True)
input_bicubic_y_image = util.resize_image_by_pil(
input_y_image,
self.scale,
resampling_method=self.resampling_method)
true_ycbcr_image = util.convert_rgb_to_ycbcr(true_image)
start = time.time()
output_y_image = self.do(input_y_image, input_bicubic_y_image)
end = time.time()
psnr, ssim = util.compute_psnr_and_ssim(
true_ycbcr_image[:, :, 0:1],
output_y_image,
border_size=self.psnr_calc_border_size)
loss_image = util.get_loss_image(
true_ycbcr_image[:, :, 0:1],
output_y_image,
border_size=self.psnr_calc_border_size)
output_color_image = util.convert_y_and_cbcr_to_rgb(
output_y_image, true_ycbcr_image[:, :, 1:3])
if save_output_images:
util.save_image(output_directory + file_path, true_image)
util.save_image(
output_directory + filename + "_input" + extension,
input_y_image)
util.save_image(
output_directory + filename + "_input_bicubic" + extension,
input_bicubic_y_image)
util.save_image(
output_directory + filename + "_true_y" + extension,
true_ycbcr_image[:, :, 0:1])
util.save_image(
output_directory + filename + "_result" + extension,
output_y_image)
util.save_image(
output_directory + filename + "_result_c" + extension,
output_color_image)
util.save_image(
output_directory + filename + "_loss" + extension,
loss_image)
elif true_image.shape[2] == 1 and self.channels == 1:
# for monochrome images
input_image = loader.build_input_image(true_image,
channels=self.channels,
scale=self.scale,
alignment=self.scale)
input_bicubic_y_image = util.resize_image_by_pil(
input_image,
self.scale,
resampling_method=self.resampling_method)
start = time.time()
output_image = self.do(input_image, input_bicubic_y_image)
end = time.time()
psnr, ssim = util.compute_psnr_and_ssim(
true_image,
output_image,
border_size=self.psnr_calc_border_size)
if save_output_images:
util.save_image(output_directory + file_path, true_image)
util.save_image(
output_directory + filename + "_result" + extension,
output_image)
else:
return None, None
if print_console:
print("[%s] PSNR:%f, SSIM:%f" % (filename, psnr, ssim))
elapsed_time = end - start
return psnr, ssim, elapsed_time
def do(self, input_image, bicubic_input_image=None):
h, w = input_image.shape[:2]
ch = input_image.shape[2] if len(input_image.shape) > 2 else 1
if bicubic_input_image is None:
bicubic_input_image = util.resize_image_by_pil(
input_image,
self.scale,
resampling_method=self.resampling_method)
if self.max_value != 255.0:
input_image = np.multiply(input_image, self.max_value /
255.0) # type: np.ndarray
bicubic_input_image = np.multiply(bicubic_input_image,
self.max_value /
255.0) # type: np.ndarray
options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
if self.self_ensemble > 1:
output = np.zeros([self.scale * h, self.scale * w, 1])
for i in range(self.self_ensemble):
image = util.flip(input_image, i)
bicubic_image = util.flip(bicubic_input_image, i)
if (self.dropout_rate < 1):
y = self.sess.run(
self.y_,
feed_dict={
self.x:
image.reshape(1, image.shape[0], image.shape[1],
ch),
self.x2:
bicubic_image.reshape(1,
self.scale * image.shape[0],
self.scale * image.shape[1],
ch),
self.dropout:
1.0,
self.is_training:
0
},
options=options,
run_metadata=run_metadata)
else:
y = self.sess.run(
self.y_,
feed_dict={
self.x:
image.reshape(1, image.shape[0], image.shape[1],
ch),
self.x2:
bicubic_image.reshape(1,
self.scale * image.shape[0],
self.scale * image.shape[1],
ch),
self.is_training:
0
},
options=options,
run_metadata=run_metadata)
restored = util.flip(y[0], i, invert=True)
output += restored
output /= self.self_ensemble
else:
y = self.sess.run(self.y_,
feed_dict={
self.x:
input_image.reshape(1, h, w, ch),
self.x2:
bicubic_input_image.reshape(
1, self.scale * h, self.scale * w, ch),
self.dropout:
1.0,
self.is_training:
0
},
options=options,
run_metadata=run_metadata)
output = y[0]
# function for timelining the model (tag: timeline)
fetched_timeline = timeline.Timeline(run_metadata.step_stats)
chrome_trace = fetched_timeline.generate_chrome_trace_format()
with open('./tf_log/inference_timeline.json', 'w') as f:
f.write(chrome_trace)
if self.max_value != 255.0:
hr_image = np.multiply(output, 255.0 / self.max_value)
else:
hr_image = output
return hr_image
def do_for_evaluate_with_output(self,
file_path,
output_directory,
print_console=False):
filename, extension = os.path.splitext(file_path)
output_directory += "/" + self.name + "/"
util.make_dir(output_directory)
true_image = util.set_image_alignment(
util.load_image(file_path, print_console=False), self.scale)
if true_image.shape[2] == 3 and self.channels == 3:
# for color images
input_image = util.build_input_image(true_image,
scale=self.scale,
alignment=self.scale)
input_bicubic_image = util.resize_image_by_pil(
input_image,
self.scale,
resampling_method=self.resampling_method)
output_image, spend_time = self.do(input_image) # SR
SR_y = eva.convert_rgb_to_y(output_image)
HR_y = eva.convert_rgb_to_y(true_image)
psnr_predicted = eva.PSNR(np.uint8(HR_y),
np.uint8(SR_y),
shave_border=self.psnr_calc_border_size)
ssim_predicted = eva.compute_ssim(np.squeeze(HR_y),
np.squeeze(SR_y))
mse = util.compute_mse(HR_y,
SR_y,
border_size=self.psnr_calc_border_size)
loss_image = util.get_loss_image(
HR_y, SR_y, border_size=self.psnr_calc_border_size)
util.save_image(output_directory + file_path[29:], true_image)
util.save_image(
output_directory + filename[28:] + "_input" + extension,
input_image)
util.save_image(
output_directory + filename[28:] + "_input_bicubic" +
extension, input_bicubic_image)
util.save_image(
output_directory + filename[28:] + "_sr" + extension,
output_image)
util.save_image(
output_directory + filename[28:] + "_loss" + extension,
loss_image)
elif true_image.shape[2] == 1 and self.channels == 1:
# for monochrome images
input_image = util.build_input_image(true_image,
scale=self.scale,
alignment=self.scale)
output_image, spend_time = self.do(input_image)
psnr_predicted = eva.PSNR(np.uint8(true_image),
np.uint8(output_image),
shave_border=self.psnr_calc_border_size)
ssim_predicted = eva.compute_ssim(np.squeeze(true_image),
np.squeeze(output_image))
mse = util.compute_mse(true_image,
output_image,
border_size=self.psnr_calc_border_size)
util.save_image(output_directory + file_path, true_image)
util.save_image(output_directory + filename + "_sr" + extension,
output_image)
else:
psnr_predicted = 0.0
ssim_predicted = 0.0
mse = 0.0
spend_time = 0.0
if print_console:
print("[%s] psnr:%f, ssim:%f, time:%f" %
(filename, psnr_predicted, ssim_predicted, spend_time))
return mse, psnr_predicted, ssim_predicted, spend_time
def load_and_evaluate_tflite_graph(
output_dir,
data_dir,
test_data,
model_path=os.path.join(os.getcwd(),
'model_to_freeze/converted_model.tflite')):
# https://stackoverflow.com/questions/50443411/how-to-load-a-tflite-model-in-script
# https://www.tensorflow.org/lite/convert/python_api#tensorflow_lite_python_interpreter_
output_directory = output_dir
output_directory += "/" + "tflite" + "/"
util.make_dir(output_directory)
test_filepaths = util.get_files_in_directory(data_dir + "/" + test_data)
total_psnr = total_ssim = total_time = 0
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path=model_path)
# interpreter = tf.contrib.lite.Interpreter(model_path=model_path)
interpreter.allocate_tensors()
# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
for file_path in test_filepaths:
# split filename from extension
filename, extension = os.path.splitext(file_path)
# prepare true image
true_image = util.set_image_alignment(
util.load_image(file_path, print_console=False), FLAGS.scale)
# start the timer
if true_image.shape[2] == 3 and FLAGS.channels == 1:
# prepare input and ground truth images
input_y_image = loader.build_input_image(true_image,
channels=FLAGS.channels,
scale=FLAGS.scale,
alignment=FLAGS.scale,
convert_ycbcr=True)
input_bicubic_y_image = util.resize_image_by_pil(
input_y_image,
FLAGS.scale,
resampling_method=FLAGS.resampling_method)
true_ycbcr_image = util.convert_rgb_to_ycbcr(true_image)
# pass inputs through the model (need to recast and reshape inputs)
input_y_image_reshaped = input_y_image.astype('float32')
input_y_image_reshaped = input_y_image_reshaped.reshape(
1, input_y_image.shape[0], input_y_image.shape[1],
FLAGS.channels)
input_bicubic_y_image_reshaped = input_bicubic_y_image.astype(
'float32')
input_bicubic_y_image_reshaped = input_bicubic_y_image_reshaped.reshape(
1, input_bicubic_y_image.shape[0],
input_bicubic_y_image.shape[1], FLAGS.channels)
interpreter.set_tensor(input_details[0]['index'],
input_y_image_reshaped) # pass x
interpreter.set_tensor(input_details[1]['index'],
input_bicubic_y_image_reshaped) # pass x2
start = time.time()
interpreter.invoke()
end = time.time()
output_y_image = interpreter.get_tensor(
output_details[0]['index']) # get y
# resize the output into an image
output_y_image = output_y_image.reshape(output_y_image.shape[1],
output_y_image.shape[2],
FLAGS.channels)
# calculate psnr and ssim for the output
psnr, ssim = util.compute_psnr_and_ssim(
true_ycbcr_image[:, :, 0:1],
output_y_image,
border_size=FLAGS.psnr_calc_border_size)
# get the loss image
loss_image = util.get_loss_image(
true_ycbcr_image[:, :, 0:1],
output_y_image,
border_size=FLAGS.psnr_calc_border_size)
# get output color image
output_color_image = util.convert_y_and_cbcr_to_rgb(
output_y_image, true_ycbcr_image[:, :, 1:3])
# save all images
util.save_image(output_directory + file_path, true_image)
util.save_image(output_directory + filename + "_input" + extension,
input_y_image)
util.save_image(
output_directory + filename + "_input_bicubic" + extension,
input_bicubic_y_image)
util.save_image(
output_directory + filename + "_true_y" + extension,
true_ycbcr_image[:, :, 0:1])
util.save_image(
output_directory + filename + "_result" + extension,
output_y_image)
util.save_image(
output_directory + filename + "_result_c" + extension,
output_color_image)
util.save_image(output_directory + filename + "_loss" + extension,
loss_image)
elapsed_time = end - start
total_psnr += psnr
total_ssim += ssim
total_time += elapsed_time
testSize = len(test_filepaths)
print("Model Average [%s] PSNR:%f, SSIM:%f, Elapsed Time:%f" %
(test_data, total_psnr / testSize, total_ssim / testSize,
total_time / testSize))
