model.py

import argparse
import glob
import sys

from absl import app
from absl.flags import argparse_flags
import tensorflow.compat.v1 as tf
import numpy as np
from PIL import Image

import tensorflow_compression as tfc

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 = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="gdn_0")),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="gdn_1")),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_2", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="gdn_2")),
            tfc.SignalConv2D(
                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 = [
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_0", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="igdn_0", inverse=True)),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="igdn_1", inverse=True)),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_2", corr=False, strides_up=2,
                padding="same_zeros", use_bias=True,
                activation=tfc.GDN(name="igdn_2", inverse=True)),
            tfc.SignalConv2D(
                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 = [
            tfc.SignalConv2D(
                self.num_filters, (3, 3), name="layer_0", corr=True, strides_down=1,
                padding="same_zeros", use_bias=True,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                self.num_filters, (5, 5), name="layer_1", corr=True, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tf.nn.relu),
            tfc.SignalConv2D(
                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 = [
            tfc.SignalConv2D(
                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),
            tfc.SignalConv2D(
                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),
            tfc.SignalConv2D(
                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


class FTransform(tf.keras.layers.Layer):
    """F function"""

    def __init__(self, num_filters, *args, **kwargs):
        self.num_filters = num_filters
        super(FTransform, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        self._layers = [
            tfc.SignalConv2D(
                self.num_filters, (3, 3), name="layer_0", corr=False, strides_up=1,
                padding="same_zeros", use_bias=True,
                activation=tf.nn.leaky_relu),
            tfc.SignalConv2D(
                self.num_filters, (3, 3), name="layer_1", corr=False, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=tf.nn.leaky_relu),
            tfc.SignalConv2D(
                self.num_filters, (3, 3), name="layer_2", corr=False, strides_down=2,
                padding="same_zeros", use_bias=True,
                activation=None),
            tfc.SignalConv2D(
                self.num_filters * 2, (1, 1), name="layer_3", corr=False, strides_down=1,
                padding="same_zeros", use_bias=True,
                activation=None),

        ]
        super(FTransform, self).build(input_shape)

    def call(self, tensor):
        for layer in self._layers:
            tensor = layer(tensor)
        return tensor


def train(args):
    """
    Train the model
    :param args:
    :return:
    """

    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 = tfc.EntropyBottleneck()
    entropy_bottleneck1 = tfc.EntropyBottleneck()
    ftransform = FTransform(args.num_filters)
    # f_transform = FTransform(args.num_filters)

    # Build autoencoder and hyperprior.
    y = analysis_transform(x)
    y_prime, y_prime_likelihoods = entropy_bottleneck1(y, training=True)
    z = hyper_analysis_transform(abs(y_prime))
    print(tf.shape(z))
    z_tilde, z_likelihoods = entropy_bottleneck(z, training=True)
    print(tf.shape(z_tilde))
    print("=====", tf.shape(y))
    y_shape = tf.shape(y)
    c_prime = hyper_synthesis_transform(z_tilde)
    c_prime = c_prime[:, :y_shape[1], :y_shape[2], :]

    mean, sigma = get_sigma_mu(y_prime, c_prime, ftransform)
    #
    scale_table = np.exp(np.linspace(
        np.log(SCALES_MIN), np.log(SCALES_MAX), SCALES_LEVELS))
    conditional_bottleneck = tfc.GaussianConditional(sigma, scale_table, mean=mean)
    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])

    aux1_optimizer = tf.train.AdamOptimizer(learning_rate=1e-3)
    aux1_step = aux1_optimizer.minimize(entropy_bottleneck1.losses[0])

    # train_op = tf.group(main_step, aux_step, entropy_bottleneck.updates[0])
    train_op = tf.group(main_step, aux_step, entropy_bottleneck.updates[0], aux1_step, entropy_bottleneck1.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 compress(args):
    """
    compress an image
    :param args:
    :return:
    """

    img = Image.open(args.input_file)
    w, h = img.size

    # Load input image and add batch dimension.
    x = read_png(args.input_file)
    x = tf.expand_dims(x, 0)
    x.set_shape([1, h, w, 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()
    entropy_bottleneck1 = tfc.EntropyBottleneck()
    ftransform = FTransform(args.num_filters)

    # Transform and compress the image.
    y = analysis_transform(x)
    y_shape = tf.shape(y)
    y_prime, y_prime_likelihoods = entropy_bottleneck1(y, training=False)
    z = hyper_analysis_transform(abs(y_prime))
    z_hat, z_likelihoods = entropy_bottleneck(z, training=False)
    c_prime = hyper_synthesis_transform(z_hat)
    c_prime = c_prime[:, :y_shape[1], :y_shape[2], :]
    mean, sigma = get_sigma_mu(y_prime, c_prime, ftransform)

    scale_table = np.exp(np.linspace(
        np.log(SCALES_MIN), np.log(SCALES_MAX), SCALES_LEVELS))
    conditional_bottleneck = tfc.GaussianConditional(sigma, scale_table, mean=mean)
    side_string = entropy_bottleneck.compress(z)
    string = conditional_bottleneck.compress(y)
    y_string = entropy_bottleneck1.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)  # 将每个维度控制在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], y_string]
        arrays = sess.run(tensors)
        print(sess.run([tf.shape(sigma), tf.shape(y), y_hat]))

        # 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, msssitm, 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 decompress(args):
    """
    decompress an image
    :param args:
    :return:
    """

    # Read the shape information and compressed string from the binary file.
    string = tf.placeholder(tf.string, [1])
    # str = tf.placeholder(tf.placeholder())
    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])
    y_string = tf.placeholder(tf.string, [1])
    with open(args.input_file, "rb") as f:
        packed = tfc.PackedTensors(f.read())
    tensors = [string, side_string, x_shape, y_shape, z_shape, y_string]
    arrays = packed.unpack(tensors)
    print(arrays)

    # Instantiate model.
    synthesis_transform = SynthesisTransform(args.num_filters)
    hyper_synthesis_transform = HyperSynthesisTransform(args.num_filters)
    entropy_bottleneck = tfc.EntropyBottleneck(dtype=tf.float32)
    entropy_bottleneck1 = tfc.EntropyBottleneck(dtype=tf.float32)
    ftransform = FTransform(args.num_filters)

    # 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)
    y_shape = tf.concat([y_shape, [args.num_filters]], axis=0)
    y_prime = entropy_bottleneck1.decompress(
        y_string, y_shape, channels=args.num_filters)
    y_prime.set_shape([1, arrays[3][0], arrays[3][1], args.num_filters])
    # y_prime_shape = tf.constant([1, arrays[3][0], arrays[3][1], args.num_filters])
    # print(y_prime_shape)
    c_prime = hyper_synthesis_transform(z_hat)
    c_prime = c_prime[:, :y_shape[0], :y_shape[1], :]
    print(tf.shape(c_prime))
    # sub_y_shape = [1, arrays[3][0], arrays[3][1], 192]
    # print("===============",sub_y_shape)
    # temp = tf.zeros(sub_y_shape, dtype=float)
    # print("===============", tf.shape(temp))
    mean_tensor, sigma_tensor = get_sigma_mu(y_prime, c_prime, ftransform)
    scale_table = np.exp(np.linspace(
        np.log(SCALES_MIN), np.log(SCALES_MAX), SCALES_LEVELS))
    conditional_bottleneck = tfc.GaussianConditional(
        sigma_tensor, scale_table, dtype=tf.float32, mean=mean_tensor)
    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 extractor_prime(padded_c_prime, h_idx, w_idx):  # with TOP N dimensions
    return padded_c_prime[:, h_idx:h_idx + 4, w_idx:w_idx + 4, :]


def extractor_doubleprime(padded_y_hat, h_idx, w_idx):  # with TOP N dimensions
    # Masking has no effect on decoding, because unknown variables are already set to zeros.
    # Therefore, the masking can be skipped in the case of decoding.
    # We just leave it here just to maintain a simple structure of code.
    # p_y_hat = padded_y_hat.transpose(0,3,1,2)
    p_y_hat = tf.transpose(padded_y_hat, [0, 3, 1, 2])
    mask = [[1, 1, 1, 1],
            [1, 1, 1, 1],
            [1, 1, 1, 1],
            [1, 1, 0, 0]]
    result = tf.math.multiply(p_y_hat[:, :, h_idx:h_idx + 4, w_idx:w_idx + 4], mask)
    # result = result.transpose(0,2,3,1)
    result = tf.transpose(result, [0, 2, 3, 1])
    return result


def get_sigma_mu(y_hat, c_prime, ftransform):
    print(tf.shape(y_hat))
    # print("y_shape: ", y_shape)
    # y_hat = tf.zeros_like(y_hat)
    paddings = tf.constant([[0, 0], [3, 0], [2, 1], [0, 0]])
    padded_y_hat = tf.pad(y_hat, paddings, mode='CONSTANT')
    padded_c_prime = tf.pad(c_prime, paddings, mode='CONSTANT')

    yshape = y_hat.get_shape()
    f_mean = tf.zeros_like(y_hat)
    f_sigma = tf.zeros_like(y_hat)
    # f_mean = tf.zeros(y_shape)
    # f_sigma = tf.zeros(y_shape)

    for h_idx in range(y_hat.shape[1]):
        for w_idx in range(y_hat.shape[2]):
            c_prime_i = extractor_prime(padded_c_prime, h_idx, w_idx)
            c_doubleprime_i = extractor_doubleprime(padded_y_hat, h_idx, w_idx)
            concatenated_c_i = tf.concat([c_doubleprime_i, c_prime_i], 3)

            f_result = ftransform(concatenated_c_i)
            # f_result = tf.layers.dense(inputs=f_result, units=384, activation=None)
            pred_mean, pred_sigma = f_result[:, :, :, 0:192], f_result[:, :, :, 191:-1]
            # pred_mean, pred_sigma = f_result, f_result
            # pmean = tf.reshape(pred_mean,[1,1,1, pred_mean.get_shape()[3].value])
            # psigma = tf.reshape(pred_sigma,[1,1,1, pred_sigma.get_shape()[3].value])
            #
            # tensorflow  tensor的修改只能针对相同shape的tensor, 因此将1x1x1x192padding为跟f_mean 跟 f_sigma具有相同的结构
            paddings = tf.constant(
                [[0, 0], [h_idx, yshape[1].value - 1 - h_idx], [w_idx, yshape[2].value - 1 - w_idx], [0, 0]])
            pred_mean = tf.pad(pred_mean, paddings, mode='CONSTANT')
            pred_sigma = tf.pad(pred_sigma, paddings, mode='CONSTANT')
            f_mean = f_mean + pred_mean
            f_sigma = f_sigma + pred_sigma

            # a  = (pred_mean + tf.constant(255.0, shape=[pred_mean.shape[3]]))
            # flatten
            # f_mean[0,h_idx, w_idx,:] = tf.reshape(a, [-1])
            # f_mean[0,h_idx, w_idx,:] = (pred_mean + np.array(255.0).repeat(pred_mean.shape[2])).flatten()
            # f_sigma[0,h_idx, w_idx,:] = tf.make_ndarray(pred_sigma, dtype=float)
            # tf.assign(f_sigma[0,h_idx, w_idx,:], pred_sigma)

    return f_mean, f_sigma


def main(args):
    # Invoke subcommand.
    if args.command == "train":
        train(args)
    elif args.command == "compress":
        if not args.output_file:
            args.output_file = args.input_file + ".tfci"
        compress(args)
    elif args.command == "decompress":
        if not args.output_file:
            args.output_file = args.input_file + ".png"
        decompress(args)

if __name__ == "__main__":
    ## compress e.g. setting the parse : --verbose --checkpoint_dir train compress  pictures/image4.jpg compressed.tfci
    ## decompress e.g setting the parse :
    app.run(main, flags_parser=parse_args)