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))