code: #%%
  """
  bmshj2018
  
  """
  
  import argparse
  import glob
  import sys
  
  from absl import app
  from absl.flags import argparse_flags
  import numpy as np
  import tensorflow.compat.v1 as tf
  
  import tensorflow_compression as tfc
  from dynamic import *
  
  
  SCALES_MIN = 0.11
  SCALES_MAX = 256
  SCALES_LEVELS = 64
  
  
  def read_png(filename):
    """Loads a PNG image file."""
    string = tf.read_file(filename)
    image = tf.image.decode_image(string, channels=3)
    image = tf.cast(image, tf.float32)
    image /= 255
    return image
  
  
  def quantize_image(image):
    image = tf.round(image * 255)
    image = tf.saturate_cast(image, tf.uint8)
    return image
  
  
  def write_png(filename, image):
    """Saves an image to a PNG file."""
    image = quantize_image(image)
    string = tf.image.encode_png(image)
    return tf.write_file(filename, string)
  
  
  class AnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform."""
  
    def __init__(self, num_filters, *args, **kwargs):
      self.num_filters = num_filters
      super(AnalysisTransform, self).__init__(*args, **kwargs)
  
    def build(self, input_shape):
      self._layers = [
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_0", corr=True, strides_down=2,
              padding="same_zeros", use_bias=True,
              activation=DynamicGDN(name="gdn_0")),
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_1", corr=True, strides_down=2,
              padding="same_zeros", use_bias=True,
              activation=DynamicGDN(name="gdn_1")),
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_2", corr=True, strides_down=2,
              padding="same_zeros", use_bias=True,
              activation=DynamicGDN(name="gdn_2")),
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_3", corr=True, strides_down=2,
              padding="same_zeros", use_bias=True,
              activation=None),
      ]
      super(AnalysisTransform, self).build(input_shape)
  
    def call(self, tensor):
      for layer in self._layers:
        tensor = layer(tensor)
      return tensor
  
  
  class SynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform."""
  
    def __init__(self, num_filters, *args, **kwargs):
      self.num_filters = num_filters
      super(SynthesisTransform, self).__init__(*args, **kwargs)
  
    def build(self, input_shape):
      self._layers = [
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
              padding="same_zeros", use_bias=True,
              activation=DynamicGDN(name="igdn_0", inverse=True)),
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_1", corr=False, strides_up=2,
              padding="same_zeros", use_bias=True,
              activation=DynamicGDN(name="igdn_1", inverse=True)),
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_2", corr=False, strides_up=2,
              padding="same_zeros", use_bias=True,
              activation=DynamicGDN(name="igdn_2", inverse=True)),
          DynamicSignalConv2D(
              3, (5, 5), name="layer_3", corr=False, strides_up=2,
              padding="same_zeros", use_bias=True,
              activation=None),
      ]
      super(SynthesisTransform, self).build(input_shape)
  
    def call(self, tensor):
      for layer in self._layers:
        tensor = layer(tensor)
      return tensor
  
  
  class HyperAnalysisTransform(tf.keras.layers.Layer):
    """The analysis transform for the entropy model parameters."""
  
    def __init__(self, num_filters, *args, **kwargs):
      self.num_filters = num_filters
      super(HyperAnalysisTransform, self).__init__(*args, **kwargs)
  
    def build(self, input_shape):
      self._layers = [
          DynamicSignalConv2D(
              self.num_filters, (3, 3), name="layer_0", corr=True, strides_down=1,
              padding="same_zeros", use_bias=True,
              activation=tf.nn.relu),
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_1", corr=True, strides_down=2,
              padding="same_zeros", use_bias=True,
              activation=tf.nn.relu),
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_2", corr=True, strides_down=2,
              padding="same_zeros", use_bias=False,
              activation=None),
      ]
      super(HyperAnalysisTransform, self).build(input_shape)
  
    def call(self, tensor):
      for layer in self._layers:
        tensor = layer(tensor)
      return tensor
  
  
  class HyperSynthesisTransform(tf.keras.layers.Layer):
    """The synthesis transform for the entropy model parameters."""
  
    def __init__(self, num_filters, *args, **kwargs):
      self.num_filters = num_filters
      super(HyperSynthesisTransform, self).__init__(*args, **kwargs)
  
    def build(self, input_shape):
      self._layers = [
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
              padding="same_zeros", use_bias=True, kernel_parameterizer=None,
              activation=tf.nn.relu),
          DynamicSignalConv2D(
              self.num_filters, (5, 5), name="layer_1", corr=False, strides_up=2,
              padding="same_zeros", use_bias=True, kernel_parameterizer=None,
              activation=tf.nn.relu),
          DynamicSignalConv2D(
              self.num_filters, (3, 3), name="layer_2", corr=False, strides_up=1,
              padding="same_zeros", use_bias=True, kernel_parameterizer=None,
              activation=None),
      ]
      super(HyperSynthesisTransform, self).build(input_shape)
  
    def call(self, tensor):
      for layer in self._layers:
        tensor = layer(tensor)
      return tensor
  
  
  def train(args):
    """Trains the model."""
  
    if args.verbose:
      tf.logging.set_verbosity(tf.logging.INFO)
  
    # Create input data pipeline.
    with tf.device("/cpu:0"):
      train_files = glob.glob(args.train_glob)
      if not train_files:
        raise RuntimeError(
            "No training images found with glob '{}'.".format(args.train_glob))
      train_dataset = tf.data.Dataset.from_tensor_slices(train_files)
      train_dataset = train_dataset.shuffle(buffer_size=len(train_files)).repeat()
      train_dataset = train_dataset.map(
          read_png, num_parallel_calls=args.preprocess_threads)
      train_dataset = train_dataset.map(
          lambda x: tf.random_crop(x, (args.patchsize, args.patchsize, 3)))
      train_dataset = train_dataset.batch(args.batchsize)
      train_dataset = train_dataset.prefetch(32)
  
    num_pixels = args.batchsize * args.patchsize ** 2
  
    # Get training patch from dataset.
    x = train_dataset.make_one_shot_iterator().get_next()
  
    # Instantiate model.
    analysis_transform = AnalysisTransform(args.num_filters)
    synthesis_transform = SynthesisTransform(args.num_filters)
    hyper_analysis_transform = HyperAnalysisTransform(args.num_filters)
    hyper_synthesis_transform = HyperSynthesisTransform(args.num_filters)
    entropy_bottleneck = DynamicEntropyBottleneck()
  
    # Build autoencoder and hyperprior.
    y = analysis_transform(x)
    z = hyper_analysis_transform(abs(y))
    z_tilde, z_likelihoods = entropy_bottleneck(z, training=True)
    sigma = hyper_synthesis_transform(z_tilde)
    scale_table = np.exp(np.linspace(
        np.log(SCALES_MIN), np.log(SCALES_MAX), SCALES_LEVELS))
    conditional_bottleneck = tfc.GaussianConditional(sigma, scale_table)
    y_tilde, y_likelihoods = conditional_bottleneck(y, training=True)
    x_tilde = synthesis_transform(y_tilde)
  
    # Total number of bits divided by number of pixels.
    train_bpp = (tf.reduce_sum(tf.log(y_likelihoods)) +
                 tf.reduce_sum(tf.log(z_likelihoods))) / (-np.log(2) * num_pixels)
  
    # Mean squared error across pixels.
    train_mse = tf.reduce_mean(tf.squared_difference(x, x_tilde))
    # Multiply by 255^2 to correct for rescaling.
    train_mse *= 255 ** 2
  
    # The rate-distortion cost.
    train_loss = args.lmbda * train_mse + train_bpp
  
    # Minimize loss and auxiliary loss, and execute update op.
    step = tf.train.create_global_step()
    main_optimizer = tf.train.AdamOptimizer(learning_rate=1e-4)
    main_step = main_optimizer.minimize(train_loss, global_step=step)
  
    aux_optimizer = tf.train.AdamOptimizer(learning_rate=1e-3)
    aux_step = aux_optimizer.minimize(entropy_bottleneck.losses[0])
  
    train_op = tf.group(main_step, aux_step, entropy_bottleneck.updates[0])
  
    tf.summary.scalar("loss", train_loss)
    tf.summary.scalar("bpp", train_bpp)
    tf.summary.scalar("mse", train_mse)
  
    tf.summary.image("original", quantize_image(x))
    tf.summary.image("reconstruction", quantize_image(x_tilde))
  
    hooks = [
        tf.train.StopAtStepHook(last_step=args.last_step),
        tf.train.NanTensorHook(train_loss),
    ]
    with tf.train.MonitoredTrainingSession(
        hooks=hooks, checkpoint_dir=args.checkpoint_dir,
        save_checkpoint_secs=300, save_summaries_secs=60) as sess:
      while not sess.should_stop():
        sess.run(train_op)
  
  
  def test_train(args):
    """Trains the model."""
  
    if args.verbose:
      tf.logging.set_verbosity(tf.logging.INFO)
  
    # Create input data pipeline.
    with tf.device("/cpu:0"):
      train_files = glob.glob(args.train_glob)
      if not train_files:
        raise RuntimeError(
            "No training images found with glob '{}'.".format(args.train_glob))
      train_dataset = tf.data.Dataset.from_tensor_slices(train_files)
      train_dataset = train_dataset.shuffle(buffer_size=len(train_files)).repeat()
      train_dataset = train_dataset.map(
          read_png, num_parallel_calls=args.preprocess_threads)
      train_dataset = train_dataset.map(
          lambda x: tf.random_crop(x, (args.patchsize, args.patchsize, 3)))
      train_dataset = train_dataset.batch(args.batchsize)
      train_dataset = train_dataset.prefetch(32)
  
    num_pixels = args.batchsize * args.patchsize ** 2
  
    # Get training patch from dataset.
    x = train_dataset.make_one_shot_iterator().get_next()
  
    # Instantiate model.
    analysis_transform = AnalysisTransform(args.num_filters)
    synthesis_transform = SynthesisTransform(args.num_filters)
    hyper_analysis_transform = HyperAnalysisTransform(args.num_filters)
    hyper_synthesis_transform = HyperSynthesisTransform(args.num_filters)
    entropy_bottleneck = DynamicEntropyBottleneck(name="entropy_bottleneck")
  
    # Build autoencoder and hyperprior.
    y = analysis_transform(x)
    z = hyper_analysis_transform(abs(y))
    z_tilde, z_likelihoods = entropy_bottleneck(z, training=True)
    sigma = hyper_synthesis_transform(z_tilde)
    scale_table = np.exp(np.linspace(
        np.log(SCALES_MIN), np.log(SCALES_MAX), SCALES_LEVELS))
    conditional_bottleneck = tfc.GaussianConditional(sigma, scale_table)
    y_tilde, y_likelihoods = conditional_bottleneck(y, training=True)
    x_tilde = synthesis_transform(y_tilde)
  
    with tf.Session() as sess:
      latest = tf.train.latest_checkpoint(checkpoint_dir="./tfc256-05")
      tf.train.Saver().restore(sess, save_path=latest)
  
    active_0 = 256
    x_tilde_0 = synthesis_transform(y_tilde[:,:,:,:active_0])
    train_bpp_0 = (tf.reduce_sum(tf.log(y_likelihoods[:,:,:,:active_0])) +
                 tf.reduce_sum(tf.log(z_likelihoods))) / (-np.log(2) * num_pixels)
    train_mse_0 = tf.reduce_mean(tf.squared_difference(x, x_tilde_0)) * (255**2)
  
    active_1 = 248
    x_tilde_1 = synthesis_transform(y_tilde[:,:,:,:active_1])
    train_bpp_1 = (tf.reduce_sum(tf.log(y_likelihoods[:,:,:,:active_1])) +
                 tf.reduce_sum(tf.log(z_likelihoods))) / (-np.log(2) * num_pixels)
    train_mse_1 = tf.reduce_mean(tf.squared_difference(x, x_tilde_1)) * (255**2)
  
    def RateOfWidth(W):
      return 0.0267 * np.exp(0.0178*W)
  
    # The rate-distortion cost.
    train_loss = train_mse_0 + train_mse_1 \
                  + 1000*tf.squared_difference(train_bpp_0, RateOfWidth(active_0)) \
                  + 1000*tf.squared_difference(train_bpp_1, RateOfWidth(active_1)) 
  
    # Minimize loss and auxiliary loss, and execute update op.
    step = tf.train.create_global_step()
    main_optimizer = tf.train.AdamOptimizer(learning_rate=1e-4)
    main_step = main_optimizer.minimize(train_loss, global_step=step)
  
    aux_optimizer = tf.train.AdamOptimizer(learning_rate=1e-3)
    aux_step = aux_optimizer.minimize(entropy_bottleneck.losses[0])
  
    train_op = tf.group(main_step, aux_step, entropy_bottleneck.updates[0])
  
    tf.summary.scalar("loss", train_loss)
    tf.summary.scalar("bpp", train_bpp_1)
    tf.summary.scalar("mse", train_mse_1)
  
    tf.summary.image("original", quantize_image(x))
    tf.summary.image("reconstruction", quantize_image(x_tilde_1))
  
    hooks = [
        tf.train.StopAtStepHook(last_step=args.last_step),
        tf.train.NanTensorHook(train_loss),
    ]
    with tf.train.MonitoredTrainingSession(
        hooks=hooks, checkpoint_dir=args.checkpoint_dir,
        save_checkpoint_secs=300, save_summaries_secs=60) as sess:
      while not sess.should_stop():
        sess.run(train_op)
  
  
  
  
  def compress(args):
    """Compresses an image."""
  
    # Load input image and add batch dimension.
    x = read_png(args.input_file)
    x = tf.expand_dims(x, 0)
    x.set_shape([1, None, None, 3])
    x_shape = tf.shape(x)
  
    # Instantiate model.
    analysis_transform = AnalysisTransform(args.num_filters)
    synthesis_transform = SynthesisTransform(args.num_filters)
    hyper_analysis_transform = HyperAnalysisTransform(args.num_filters)
    hyper_synthesis_transform = HyperSynthesisTransform(args.num_filters)
    entropy_bottleneck = tfc.EntropyBottleneck()
  
    # Transform and compress the image.
    y = analysis_transform(x)
    y_shape = tf.shape(y)
    z = hyper_analysis_transform(abs(y))
    z_hat, z_likelihoods = entropy_bottleneck(z, training=False)
    sigma = hyper_synthesis_transform(z_hat)
    sigma = sigma[:, :y_shape[1], :y_shape[2], :]
    scale_table = np.exp(np.linspace(
        np.log(SCALES_MIN), np.log(SCALES_MAX), SCALES_LEVELS))
    conditional_bottleneck = tfc.GaussianConditional(sigma, scale_table)
    side_string = entropy_bottleneck.compress(z)
    string = conditional_bottleneck.compress(y)
  
    # Transform the quantized image back (if requested).
    y_hat, y_likelihoods = conditional_bottleneck(y, training=False)
    x_hat = synthesis_transform(y_hat)
    x_hat = x_hat[:, :x_shape[1], :x_shape[2], :]
  
    num_pixels = tf.cast(tf.reduce_prod(tf.shape(x)[:-1]), dtype=tf.float32)
  
    # Total number of bits divided by number of pixels.
    eval_bpp = (tf.reduce_sum(tf.log(y_likelihoods)) +
                tf.reduce_sum(tf.log(z_likelihoods))) / (-np.log(2) * num_pixels)
  
    # Bring both images back to 0..255 range.
    x *= 255
    x_hat = tf.clip_by_value(x_hat, 0, 1)
    x_hat = tf.round(x_hat * 255)
  
    mse = tf.reduce_mean(tf.squared_difference(x, x_hat))
    psnr = tf.squeeze(tf.image.psnr(x_hat, x, 255))
    msssim = tf.squeeze(tf.image.ssim_multiscale(x_hat, x, 255))
  
    with tf.Session() as sess:
      # Load the latest model checkpoint, get the compressed string and the tensor
      # shapes.
      latest = tf.train.latest_checkpoint(checkpoint_dir=args.checkpoint_dir)
      tf.train.Saver().restore(sess, save_path=latest)
      tensors = [string, side_string,
                 tf.shape(x)[1:-1], tf.shape(y)[1:-1], tf.shape(z)[1:-1]]
      arrays = sess.run(tensors)
  
      # Write a binary file with the shape information and the compressed string.
      packed = tfc.PackedTensors()
      packed.pack(tensors, arrays)
      with open(args.output_file, "wb") as f:
        f.write(packed.string)
  
      # If requested, transform the quantized image back and measure performance.
      if args.verbose:
        eval_bpp, mse, psnr, msssim, num_pixels = sess.run(
            [eval_bpp, mse, psnr, msssim, num_pixels])
  
        # The actual bits per pixel including overhead.
        bpp = len(packed.string) * 8 / num_pixels
  
        print("Mean squared error: {:0.4f}".format(mse))
        print("PSNR (dB): {:0.2f}".format(psnr))
        print("Multiscale SSIM: {:0.4f}".format(msssim))
        print("Multiscale SSIM (dB): {:0.2f}".format(-10 * np.log10(1 - msssim)))
        print("Information content in bpp: {:0.4f}".format(eval_bpp))
        print("Actual bits per pixel: {:0.4f}".format(bpp))
  
  
  def get_uninitialized_variables(sess):
    global_vars = tf.global_variables()
  
    # print([str(i.name) for i in global_vars])
  
    is_not_initialized = sess.run([tf.is_variable_initialized(var) for var in global_vars])
    not_initialized_vars = [v for (v, f) in zip(global_vars, is_not_initialized) if not f]
    print([str(i.name) for i in not_initialized_vars])
    return not_initialized_vars
  
  
  class DyTFC():
    def __init__(self, num_filters):
      self.num_filters = num_filters
      self.analysis_transform = AnalysisTransform(num_filters)
      self.synthesis_transform = SynthesisTransform(num_filters)
      self.hyper_analysis_transform = HyperAnalysisTransform(num_filters)
      self.hyper_synthesis_transform = HyperSynthesisTransform(num_filters)
      self.entropy_bottleneck = DynamicEntropyBottleneck(name="entropy_bottleneck")
      
    def build(self, x):
      self.x = x
      self.x_shape = tf.shape(self.x)
  
      # Transform and compress the image.
      self.y = self.analysis_transform(self.x)
      self.y_shape = tf.shape(self.y)
      self.z = self.hyper_analysis_transform(abs(self.y))
      self.z_shape = tf.shape(self.z)
      self.z_hat, self.z_likelihoods = self.entropy_bottleneck(self.z, training=False)
      sigma = self.hyper_synthesis_transform(self.z_hat)
      self.sigma = sigma[:, :self.y_shape[1], :self.y_shape[2], :]
      self.scale_table = np.exp(np.linspace(
          np.log(SCALES_MIN), np.log(SCALES_MAX), SCALES_LEVELS))
      self.conditional_bottleneck = DynamicGaussianConditional(self.sigma, self.scale_table, name="gaussian_conditional")
      self.side_string = self.entropy_bottleneck.compress(self.z)
      self.string = self.conditional_bottleneck.compress(self.y)
  
      # Transform the quantized image back (if requested).
      self.y_hat, self.y_likelihoods = self.conditional_bottleneck(self.y, training=False)
      self.x_hat = self.synthesis_transform(self.y_hat)
      self.x_hat = self.x_hat[:, :self.x_shape[1], :self.x_shape[2], :]
  
      self.num_pixels = tf.cast(tf.reduce_prod(tf.shape(self.x)[:-1]), dtype=tf.float32)
  
      # Total number of bits divided by number of pixels.
      self.eval_bpp = (tf.reduce_sum(tf.log(self.y_likelihoods)) +
                  tf.reduce_sum(tf.log(self.z_likelihoods))) / (-np.log(2) * self.num_pixels)
  
      x = self.x * 255
      x_hat = tf.clip_by_value(self.x_hat, 0, 1)
      x_hat = tf.round(x_hat * 255)
  
      self.mse = tf.reduce_mean(tf.squared_difference(x, x_hat))
      self.psnr = tf.squeeze(tf.image.psnr(x_hat, x, 255))
      self.msssim = tf.squeeze(tf.image.ssim_multiscale(x_hat, x, 255))
  
      self.vst = {v.name:v for v in tf.global_variables()}
  
    def _reorg(self, sess, trans, active_out_filters, sort_in, sort_out, flag=None):
      layers = trans._layers
      sorted_idx = sort_in
      for layer in layers[:-1]:
        layer.active_out_filters = active_out_filters
        sorted_idx = layer.sort_filter(sess, self.vst, sorted_idx, True)
      if flag is not "tail":
        layers[-1].active_out_filters = active_out_filters
      sorted_idx = layers[-1].sort_filter(sess, self.vst, sorted_idx, sort_out)
      return sorted_idx
  
    def reorg(self, sess, active):
      y_sorted_idx = self._reorg(sess, self.analysis_transform,        active, False, True, "head")
      # print(sess.run(y_sorted_idx))
      self._reorg(               sess, self.synthesis_transform,       active, y_sorted_idx, False, "tail")
      z_sorted_idx = self._reorg(sess, self.hyper_analysis_transform,  active, y_sorted_idx, True, "body")
      self.entropy_bottleneck.sort_weight(sess, self.vst, z_sorted_idx)
      self.entropy_bottleneck.input_spec = tf.keras.layers.InputSpec(ndim=4, axes={3: active})
      sorted_idx = self._reorg(  sess, self.hyper_synthesis_transform,   active, z_sorted_idx, y_sorted_idx, "body")
  
    def _active(self, trans, active_out_filters, flag=None):
      layers = trans._layers
      for layer in layers[:-1]:
        layer.active_out_filters = active_out_filters
      if flag is not "tail":
        layers[-1].active_out_filters = active_out_filters
  
    def set_active(self, active):
      self._active(self.analysis_transform, active, "head")
      self._active(self.synthesis_transform, active, "tail")
      self._active(self.hyper_analysis_transform, active, "body")
      self._active(self.hyper_synthesis_transform, active, "body")
      self.entropy_bottleneck.input_spec = tf.keras.layers.InputSpec(ndim=4, axes={3: active})
      self.conditional_bottleneck.input_spec = tf.keras.layers.InputSpec(ndim=4, axes={3: active})
  
  
  def test_compress(args):
    """Compresses an image."""
  
    # Load input image and add batch dimension.
    x = read_png(args.input_file)
    x = tf.expand_dims(x, 0)
    x.set_shape([1, None, None, 3])
    x_shape = tf.shape(x)
  
    net = DyTFC(192)
    net.build(x)
  
    sess =  tf.Session()
    latest = tf.train.latest_checkpoint(checkpoint_dir=args.checkpoint_dir)
    tf.train.Saver().restore(sess, save_path=latest)
    print(sess.run( tf.reduce_sum(tf.log(net.y_likelihoods), axis=(0,1,2)) / (-np.log(2) * net.num_pixels)) )
    return
  
    #vnames = ['gaussian_conditional/quantized_cdf:0', 'gaussian_conditional/cdf_length:0']
    #old_cb_weights = net.conditional_bottleneck.get_weights()
  
    #print(old_cb_weights)
    #net.set_active(192)
    #net.build(x)
    #sess.run(tf.variables_initializer(get_uninitialized_variables(sess)))
    #sess.run(tf.variables_initializer([net.vst[name] for name in vnames]))
    #net.conditional_bottleneck.set_weights(old_cb_weights)
  
    #
    #tf.train.Saver().save(sess,"./sort128/model.ckpt")  
  
    tensors = [net.string, net.side_string,
                net.x_shape[1:-1], net.y_shape[1:-1], net.z_shape[1:-1]]
    
    arrays = sess.run(tensors)
  
    # Write a binary file with the shape information and the compressed string.
    packed = tfc.PackedTensors()
    packed.pack(tensors, arrays)
    with open(args.output_file, "wb") as f:
      f.write(packed.string)
  
    # If requested, transform the quantized image back and measure performance.
    if args.verbose:
      eval_bpp, mse, psnr, msssim, num_pixels = sess.run(
          [net.eval_bpp, net.mse, net.psnr, net.msssim, net.num_pixels])
  
      # The actual bits per pixel including overhead.
      bpp = len(packed.string) * 8 / num_pixels
  
      print("Mean squared error: {:0.4f}".format(mse))
      print("PSNR (dB): {:0.2f}".format(psnr))
      print("Multiscale SSIM: {:0.4f}".format(msssim))
      print("Multiscale SSIM (dB): {:0.2f}".format(-10 * np.log10(1 - msssim)))
      print("Information content in bpp: {:0.4f}".format(eval_bpp))
      print("Actual bits per pixel: {:0.4f}".format(bpp))
    
  
  
  def test_decompress(args):
    """Decompresses an image."""
  
    # Read the shape information and compressed string from the binary file.
    string = tf.placeholder(tf.string, [1])
    side_string = tf.placeholder(tf.string, [1])
    x_shape = tf.placeholder(tf.int32, [2])
    y_shape = tf.placeholder(tf.int32, [2])
    z_shape = tf.placeholder(tf.int32, [2])
    with open(args.input_file, "rb") as f:
      packed = tfc.PackedTensors(f.read())
    tensors = [string, side_string, x_shape, y_shape, z_shape]
    arrays = packed.unpack(tensors)
  
    # Instantiate model.
    synthesis_transform = SynthesisTransform(args.num_filters)
    hyper_synthesis_transform = HyperSynthesisTransform(args.num_filters)
    entropy_bottleneck = tfc.EntropyBottleneck(dtype=tf.float32)
  
    # Decompress and transform the image back.
    z_shape = tf.concat([z_shape, [args.num_filters]], axis=0)
    z_hat = entropy_bottleneck.decompress(
        side_string, z_shape, channels=args.num_filters)
    sigma = hyper_synthesis_transform(z_hat)
    sigma = sigma[:, :y_shape[0], :y_shape[1], :]
    scale_table = np.exp(np.linspace(
        np.log(SCALES_MIN), np.log(SCALES_MAX), SCALES_LEVELS))
    conditional_bottleneck = tfc.GaussianConditional(
        sigma, scale_table, dtype=tf.float32)
    y_hat_all = conditional_bottleneck.decompress(string)
  
    x = read_png("kodak/kodim01.png")
    x = tf.expand_dims(x, 0)
    x.set_shape([1, None, None, 3])
    x_shape = tf.shape(x)
    x *= 255
  
    active = 192
    y_hat = y_hat_all[:,:,:,:active]
    x_hat = synthesis_transform(y_hat)
    x_hat = tf.clip_by_value(x_hat, 0, 1)
    x_hat = tf.round(x_hat * 255)
    mse = tf.reduce_mean(tf.squared_difference(x, x_hat))
    psnr = tf.squeeze(tf.image.psnr(x_hat, x, 255))
    msssim = tf.squeeze(tf.image.ssim_multiscale(x_hat, x, 255))
  
    #x_hat = x_hat[0, :x_shape[0], :x_shape[1], :]
    #op = write_png(args.output_file, x_hat)
  
    sess = tf.Session()
    latest = tf.train.latest_checkpoint(checkpoint_dir=args.checkpoint_dir)
    tf.train.Saver().restore(sess, save_path=latest)
    #sess.run(op, feed_dict=dict(zip(tensors, arrays)))
  
    #vmse, vpsnr, vmsssim = sess.run([mse, psnr, msssim], feed_dict=dict(zip(tensors, arrays)))
    #print(vmse, vpsnr, vmsssim)
  
    for active in range(192,0,-8):
      y_hat = y_hat_all[:,:,:,:active]
      x_hat = synthesis_transform(y_hat)
      x_hat = tf.clip_by_value(x_hat, 0, 1)
      x_hat = tf.round(x_hat * 255)
      mse = tf.reduce_mean(tf.squared_difference(x, x_hat))
      psnr = tf.squeeze(tf.image.psnr(x_hat, x, 255))
      msssim = tf.squeeze(tf.image.ssim_multiscale(x_hat, x, 255))
      vmse, vpsnr, vmsssim = sess.run([mse, psnr, msssim], feed_dict=dict(zip(tensors, arrays)))
      print(active, vmse, vpsnr, vmsssim)
  
  
  
  
  def decompress(args):
    """Decompresses an image."""
  
    # Read the shape information and compressed string from the binary file.
    string = tf.placeholder(tf.string, [1])
    side_string = tf.placeholder(tf.string, [1])
    x_shape = tf.placeholder(tf.int32, [2])
    y_shape = tf.placeholder(tf.int32, [2])
    z_shape = tf.placeholder(tf.int32, [2])
    with open(args.input_file, "rb") as f:
      packed = tfc.PackedTensors(f.read())
    tensors = [string, side_string, x_shape, y_shape, z_shape]
    arrays = packed.unpack(tensors)
  
    # Instantiate model.
    synthesis_transform = SynthesisTransform(args.num_filters)
    hyper_synthesis_transform = HyperSynthesisTransform(args.num_filters)
    entropy_bottleneck = tfc.EntropyBottleneck(dtype=tf.float32)
  
    # Decompress and transform the image back.
    z_shape = tf.concat([z_shape, [args.num_filters]], axis=0)
    z_hat = entropy_bottleneck.decompress(
        side_string, z_shape, channels=args.num_filters)
    sigma = hyper_synthesis_transform(z_hat)
    sigma = sigma[:, :y_shape[0], :y_shape[1], :]
    scale_table = np.exp(np.linspace(
        np.log(SCALES_MIN), np.log(SCALES_MAX), SCALES_LEVELS))
    conditional_bottleneck = tfc.GaussianConditional(
        sigma, scale_table, dtype=tf.float32)
    y_hat = conditional_bottleneck.decompress(string)
    x_hat = synthesis_transform(y_hat)
  
    # Remove batch dimension, and crop away any extraneous padding on the bottom
    # or right boundaries.
    x_hat = x_hat[0, :x_shape[0], :x_shape[1], :]
  
    # Write reconstructed image out as a PNG file.
    op = write_png(args.output_file, x_hat)
  
    # Load the latest model checkpoint, and perform the above actions.
    with tf.Session() as sess:
      latest = tf.train.latest_checkpoint(checkpoint_dir=args.checkpoint_dir)
      tf.train.Saver().restore(sess, save_path=latest)
      sess.run(op, feed_dict=dict(zip(tensors, arrays)))
  
  
  def parse_args(argv):
    """Parses command line arguments."""
    parser = argparse_flags.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  
    # High-level options.
    parser.add_argument(
        "--verbose", "-V", action="store_true",
        help="Report bitrate and distortion when training or compressing.")
    parser.add_argument(
        "--num_filters", type=int, default=192,
        help="Number of filters per layer.")
    parser.add_argument(
        "--checkpoint_dir", default="train",
        help="Directory where to save/load model checkpoints.")
    subparsers = parser.add_subparsers(
        title="commands", dest="command",
        help="What to do: 'train' loads training data and trains (or continues "
             "to train) a new model. 'compress' reads an image file (lossless "
             "PNG format) and writes a compressed binary file. 'decompress' "
             "reads a binary file and reconstructs the image (in PNG format). "
             "input and output filenames need to be provided for the latter "
             "two options. Invoke '<command> -h' for more information.")
  
    # 'train' subcommand.
    train_cmd = subparsers.add_parser(
        "train",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Trains (or continues to train) a new model.")
    train_cmd.add_argument(
        "--train_glob", default="images/*.png",
        help="Glob pattern identifying training data. This pattern must expand "
             "to a list of RGB images in PNG format.")
    train_cmd.add_argument(
        "--batchsize", type=int, default=8,
        help="Batch size for training.")
    train_cmd.add_argument(
        "--patchsize", type=int, default=256,
        help="Size of image patches for training.")
    train_cmd.add_argument(
        "--lambda", type=float, default=0.01, dest="lmbda",
        help="Lambda for rate-distortion tradeoff.")
    train_cmd.add_argument(
        "--last_step", type=int, default=1000000,
        help="Train up to this number of steps.")
    train_cmd.add_argument(
        "--preprocess_threads", type=int, default=16,
        help="Number of CPU threads to use for parallel decoding of training "
             "images.")
  
    # 'compress' subcommand.
    compress_cmd = subparsers.add_parser(
        "compress",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Reads a PNG file, compresses it, and writes a TFCI file.")
  
    # 'decompress' subcommand.
    decompress_cmd = subparsers.add_parser(
        "decompress",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Reads a TFCI file, reconstructs the image, and writes back "
                    "a PNG file.")
  
    # Arguments for both 'compress' and 'decompress'.
    for cmd, ext in ((compress_cmd, ".tfci"), (decompress_cmd, ".png")):
      cmd.add_argument(
          "input_file",
          help="Input filename.")
      cmd.add_argument(
          "output_file", nargs="?",
          help="Output filename (optional). If not provided, appends '{}' to "
               "the input filename.".format(ext))
  
    # Parse arguments.
    args = parser.parse_args(argv[1:])
    if args.command is None:
      parser.print_usage()
      sys.exit(2)
    return args
  
  
  def main(args):
    # Invoke subcommand.
    if args.command == "train":
      test_train(args)
    elif args.command == "compress":
      if not args.output_file:
        args.output_file = args.input_file + ".tfci"
      test_compress(args)
    elif args.command == "decompress":
      if not args.output_file:
        args.output_file = args.input_file + ".png"
      test_decompress(args)
  
  #%%
  
  #%%
  
  if __name__ == "__main__":
    app.run(main, flags_parser=parse_args)
  
apis: tensorflow.compat.v1.train.MonitoredTrainingSession tensorflow.compat.v1.train.StopAtStepHook tensorflow.compat.v1.log tensorflow.compat.v1.train.AdamOptimizer tensorflow.compat.v1.expand_dims numpy.log tensorflow.compat.v1.image.psnr tensorflow.compat.v1.shape tensorflow.compat.v1.clip_by_value tensorflow.compat.v1.logging.set_verbosity tensorflow.compat.v1.placeholder tensorflow.compat.v1.train.NanTensorHook tensorflow.compat.v1.group tensorflow.compat.v1.train.Saver tensorflow.compat.v1.train.create_global_step tensorflow.compat.v1.device tensorflow.compat.v1.read_file numpy.log10 tensorflow.compat.v1.cast tensorflow.compat.v1.saturate_cast tensorflow.compat.v1.keras.layers.InputSpec tensorflow.compat.v1.data.Dataset.from_tensor_slices tensorflow.compat.v1.Session tensorflow.compat.v1.is_variable_initialized tensorflow.compat.v1.concat tensorflow.compat.v1.image.encode_png tensorflow.compat.v1.image.decode_image tensorflow.compat.v1.write_file tensorflow.compat.v1.summary.scalar tensorflow.compat.v1.image.ssim_multiscale tensorflow.compat.v1.round tensorflow.compat.v1.train.latest_checkpoint tensorflow.compat.v1.global_variables numpy.exp tensorflow.compat.v1.squared_difference tensorflow.compat.v1.random_crop