text: import os
  import tensorflow as tf
  import numpy as np
  
  from tensorflow.python.data.experimental import AUTOTUNE
  from tensorflow import keras
  
  
  import matplotlib.pyplot as plt
  from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes
  from mpl_toolkits.axes_grid1.inset_locator import mark_inset
  import PIL
  from tensorflow.keras.preprocessing.image import load_img
  from tensorflow.keras.preprocessing.image import array_to_img
  from tensorflow.keras.preprocessing.image import img_to_array
  
  def scaling(input_image):
      input_image = input_image / 255.0
      return input_image
  
  # Use TF Ops to process.
  def process_input(input, input_size, upscale_factor):
      return tf.image.resize(input, [input_size, input_size], method="area")
  
  
  def process_target(input):
      return input
  
  
  def plot_results(img, prefix, title):
      """Plot the result with zoom-in area."""
      img_array = img_to_array(img)
      img_array = img_array.astype("float32") / 255.0
  
      # Create a new figure with a default 111 subplot.
      fig, ax = plt.subplots()
      im = ax.imshow(img_array[::-1], origin="lower")
  
      plt.title(title)
      # zoom-factor: 2.0, location: upper-left
      axins = zoomed_inset_axes(ax, 2, loc=2)
      axins.imshow(img_array[::-1], origin="lower")
  
      # Specify the limits.
      x1, x2, y1, y2 = 200, 300, 100, 200
      # Apply the x-limits.
      axins.set_xlim(x1, x2)
      # Apply the y-limits.
      axins.set_ylim(y1, y2)
  
      plt.yticks(visible=False)
      plt.xticks(visible=False)
  
      # Make the line.
      mark_inset(ax, axins, loc1=1, loc2=3, fc="none", ec="blue")
      plt.savefig(str(prefix) + "-" + title + ".png")
  
  
  def get_lowres_image(img, upscale_factor):
      """Return low-resolution image to use as model input."""
      return img.resize(
          (img.size[0] // upscale_factor, img.size[1] // upscale_factor),
          PIL.Image.BICUBIC,
      )
  
  
  def upscale_image(model, img):
      """Predict the result based on input image and restore the image as RGB."""
      y = img_to_array(img)
      y = y.astype("float32") / 255.0
  
      input = np.expand_dims(y, axis=0)
      out = model.predict(input)
  
      out_img_y = out[0]
      out_img_y *= 255.0
  
      # Restore the image in RGB color space.
      out_img_y = out_img_y.clip(0, 255)
      out_img_y = out_img_y.reshape((np.shape(out_img_y)[0], np.shape(out_img_y)[1], 3))
      out_img = PIL.Image.fromarray(np.uint8(out_img_y))
      # out_img_cb = cb.resize(out_img_y.size, PIL.Image.BICUBIC)
      # out_img_cr = cr.resize(out_img_y.size, PIL.Image.BICUBIC)
      # out_img = PIL.Image.merge("YCbCr", (out_img_y, out_img_cb, out_img_cr)).convert(
      #     "RGB"
      # )
      return out_img
  
  class ESPCNCallback(keras.callbacks.Callback):
      def __init__(self):
          super(ESPCNCallback, self).__init__()
          self.test_img = get_lowres_image(load_img(test_img_paths[0]), upscale_factor)
  
      # Store PSNR value in each epoch.
      def on_epoch_begin(self, epoch, logs=None):
          self.psnr = []
  
      def on_epoch_end(self, epoch, logs=None):
          print("Mean PSNR for epoch: %.2f" % (np.mean(self.psnr)))
          if epoch % 10 == 0:
              prediction = upscale_image(self.model, self.test_img)
              plot_results(prediction, "epoch-" + str(epoch), "prediction")
  
      def on_test_batch_end(self, batch, logs=None):
          self.psnr.append(10 * math.log10(1 / logs["loss"]))
  
  
  class DIV2K:
      def __init__(self,
                   scale=2,
                   subset='train',
                   downgrade='bicubic',
                   images_dir='div2k',
                   caches_dir='div2k/caches'):
  
          self._ntire_2018 = True
  
          _scales = [2, 3, 4, 8]
  
          if scale in _scales:
              self.scale = scale
          else:
              raise ValueError(f'scale must be in ${_scales}')
  
          if subset == 'train':
              self.image_ids = range(1, 801)
          elif subset == 'valid':
              self.image_ids = range(801, 901)
          else:
              raise ValueError("subset must be 'train' or 'valid'")
  
          _downgrades_a = ['bicubic', 'unknown']
          _downgrades_b = ['mild', 'difficult']
  
          if scale == 8 and downgrade != 'bicubic':
              raise ValueError(f'scale 8 only allowed for bicubic downgrade')
  
          if downgrade in _downgrades_b and scale != 4:
              raise ValueError(f'{downgrade} downgrade requires scale 4')
  
          if downgrade == 'bicubic' and scale == 8:
              self.downgrade = 'x8'
          elif downgrade in _downgrades_b:
              self.downgrade = downgrade
          else:
              self.downgrade = downgrade
              self._ntire_2018 = False
  
          self.subset = subset
          self.images_dir = images_dir
          self.caches_dir = caches_dir
  
          os.makedirs(images_dir, exist_ok=True)
          os.makedirs(caches_dir, exist_ok=True)
  
      def __len__(self):
          return len(self.image_ids)
  
      def dataset(self, batch_size=8, repeat_count=None, random_transform=True):
          ds = tf.data.Dataset.zip((self.lr_dataset(), self.hr_dataset()))
          if random_transform:
              ds = ds.map(lambda lr, hr: random_crop(lr, hr, scale=self.scale), num_parallel_calls=AUTOTUNE)
              ds = ds.map(random_rotate, num_parallel_calls=AUTOTUNE)
              ds = ds.map(random_flip, num_parallel_calls=AUTOTUNE)
          ds = ds.batch(batch_size)
          ds = ds.repeat(repeat_count)
          ds = ds.prefetch(buffer_size=AUTOTUNE)
          return ds
  
      def hr_dataset(self):
          if not os.path.exists(self._hr_images_dir()):
              download_archive(self._hr_images_archive(), self.images_dir, extract=True)
  
          ds = self._images_dataset(self._hr_image_files()).cache(self._hr_cache_file())
  
          if not os.path.exists(self._hr_cache_index()):
              self._populate_cache(ds, self._hr_cache_file())
  
          return ds
  
      def lr_dataset(self):
          if not os.path.exists(self._lr_images_dir()):
              download_archive(self._lr_images_archive(), self.images_dir, extract=True)
  
          ds = self._images_dataset(self._lr_image_files()).cache(self._lr_cache_file())
  
          if not os.path.exists(self._lr_cache_index()):
              self._populate_cache(ds, self._lr_cache_file())
  
          return ds
  
      def _hr_cache_file(self):
          return os.path.join(self.caches_dir, f'DIV2K_{self.subset}_HR.cache')
  
      def _lr_cache_file(self):
          return os.path.join(self.caches_dir, f'DIV2K_{self.subset}_LR_{self.downgrade}_X{self.scale}.cache')
  
      def _hr_cache_index(self):
          return f'{self._hr_cache_file()}.index'
  
      def _lr_cache_index(self):
          return f'{self._lr_cache_file()}.index'
  
      def _hr_image_files(self):
          images_dir = self._hr_images_dir()
          return [os.path.join(images_dir, f'{image_id:04}.png') for image_id in self.image_ids]
  
      def _lr_image_files(self):
          images_dir = self._lr_images_dir()
          return [os.path.join(images_dir, self._lr_image_file(image_id)) for image_id in self.image_ids]
  
      def _lr_image_file(self, image_id):
          if not self._ntire_2018 or self.scale == 8:
              return f'{image_id:04}x{self.scale}.png'
          else:
              return f'{image_id:04}x{self.scale}{self.downgrade[0]}.png'
  
      def _hr_images_dir(self):
          return os.path.join(self.images_dir, f'DIV2K_{self.subset}_HR')
  
      def _lr_images_dir(self):
          if self._ntire_2018:
              return os.path.join(self.images_dir, f'DIV2K_{self.subset}_LR_{self.downgrade}')
          else:
              return os.path.join(self.images_dir, f'DIV2K_{self.subset}_LR_{self.downgrade}', f'X{self.scale}')
  
      def _hr_images_archive(self):
          return f'DIV2K_{self.subset}_HR.zip'
  
      def _lr_images_archive(self):
          if self._ntire_2018:
              return f'DIV2K_{self.subset}_LR_{self.downgrade}.zip'
          else:
              return f'DIV2K_{self.subset}_LR_{self.downgrade}_X{self.scale}.zip'
  
      @staticmethod
      def _images_dataset(image_files):
          ds = tf.data.Dataset.from_tensor_slices(image_files)
          ds = ds.map(tf.io.read_file)
          ds = ds.map(lambda x: tf.image.decode_png(x, channels=3), num_parallel_calls=AUTOTUNE)
          return ds
  
      @staticmethod
      def _populate_cache(ds, cache_file):
          print(f'Caching decoded images in {cache_file} ...')
          for _ in ds: pass
          print(f'Cached decoded images in {cache_file}.')
  
  
  # -----------------------------------------------------------
  #  Transformations
  # -----------------------------------------------------------
  
  
  def random_crop(lr_img, hr_img, hr_crop_size=512, scale=2):
      lr_crop_size = hr_crop_size // scale
      lr_img_shape = tf.shape(lr_img)[:2]
  
      lr_w = tf.random.uniform(shape=(), maxval=lr_img_shape[1] - lr_crop_size + 1, dtype=tf.int32)
      lr_h = tf.random.uniform(shape=(), maxval=lr_img_shape[0] - lr_crop_size + 1, dtype=tf.int32)
  
      hr_w = lr_w * scale
      hr_h = lr_h * scale
  
      lr_img_cropped = lr_img[lr_h:lr_h + lr_crop_size, lr_w:lr_w + lr_crop_size]
      hr_img_cropped = hr_img[hr_h:hr_h + hr_crop_size, hr_w:hr_w + hr_crop_size]
  
      return lr_img_cropped, hr_img_cropped
  
  
  def random_flip(lr_img, hr_img):
      rn = tf.random.uniform(shape=(), maxval=1)
      return tf.cond(rn < 0.5,
                     lambda: (lr_img, hr_img),
                     lambda: (tf.image.flip_left_right(lr_img),
                              tf.image.flip_left_right(hr_img)))
  
  
  def random_rotate(lr_img, hr_img):
      rn = tf.random.uniform(shape=(), maxval=4, dtype=tf.int32)
      return tf.image.rot90(lr_img, rn), tf.image.rot90(hr_img, rn)
  
  
  # -----------------------------------------------------------
  #  IO
  # -----------------------------------------------------------
  
  
  def download_archive(file, target_dir, extract=True):
      source_url = f'http://data.vision.ee.ethz.ch/cvl/DIV2K/{file}'
      target_dir = os.path.abspath(target_dir)
      tf.keras.utils.get_file(file, source_url, cache_subdir=target_dir, extract=extract)
      os.remove(os.path.join(target_dir, file))
  
  
  def evaluate(model, dataset):
      psnr_values = []
      for data in dataset:
          lr = list(data)[0]
          hr = list(data)[1]
          lr = tf.image.resize(lr, (256, 256))
          hr = tf.image.resize(hr, (512, 512))
          sr = resolve(model, lr)
          psnr_value = psnr(hr, sr)[0]
          psnr_values.append(psnr_value)
      return tf.reduce_mean(psnr_values)
  
  
  def resolve(model, lr_batch):
      lr_batch = tf.cast(lr_batch, tf.float32)
      sr_batch = model(lr_batch)
      sr_batch = tf.clip_by_value(sr_batch, 0 ,255)
      sr_batch = tf.round(sr_batch)
      sr_batch = tf.cast(sr_batch, tf.uint8)
  
      return sr_batch
  
  def psnr(x1, x2):
      return tf.image.psnr(x1, x2, max_val=255)
  
  