Refactored project code

fixed the entropy model and improved training speed by using TFRecords format for the data set, a script is included to construct the records from any image dataset

Author: MahmoudAshraf97

LPIPS.py

@@ -0,0 +1,53 @@
+import os
+import tensorflow as tf
+import urllib.request
+
+_LPIPS_URL = "http://rail.eecs.berkeley.edu/models/lpips/net-lin_alex_v0.1.pb"
+
+
+def ensure_lpips_weights_exist(weight_path_out):
+    """Downloads weights if needed."""
+    if os.path.isfile(weight_path_out):
+        return
+    print("Downloading LPIPS weights:", _LPIPS_URL, "->", weight_path_out)
+    urllib.request.urlretrieve(_LPIPS_URL, weight_path_out)
+    if not os.path.isfile(weight_path_out):
+        raise ValueError(f"Failed to download LPIPS weights from {_LPIPS_URL} "
+                         f"to {weight_path_out}. Please manually download!")
+
+
+class LPIPSLoss(object):
+    """Calcualte LPIPS loss."""
+
+    def __init__(self, weight_path):
+        ensure_lpips_weights_exist(weight_path)
+
+        def wrap_frozen_graph(graph_def, inputs, outputs):
+            def _imports_graph_def():
+                tf.graph_util.import_graph_def(graph_def, name="")
+
+            wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])
+            import_graph = wrapped_import.graph
+            return wrapped_import.prune(
+                tf.nest.map_structure(import_graph.as_graph_element, inputs),
+                tf.nest.map_structure(import_graph.as_graph_element, outputs))
+
+        # Pack LPIPS network into a tf function
+        graph_def = tf.compat.v1.GraphDef()
+        with open(weight_path, "rb") as f:
+            graph_def.ParseFromString(f.read())
+        self._lpips_func = tf.function(
+            wrap_frozen_graph(
+                graph_def, inputs=("0:0", "1:0"), outputs="Reshape_10:0"))
+
+    def __call__(self, fake_image, real_image):
+        """Assuming inputs are in [0, 1]."""
+
+        # Move inputs to [-1, 1] and NCHW format.
+        def _transpose_to_nchw(x):
+            return tf.transpose(x, (0, 3, 1, 2))
+
+        fake_image = _transpose_to_nchw(fake_image * 2 - 1.0)
+        real_image = _transpose_to_nchw(real_image * 2 - 1.0)
+        loss = self._lpips_func(fake_image, real_image)
+        return tf.reduce_mean(loss)  # Loss is N111, take mean to get scalar.

compress.py

@@ -0,0 +1,59 @@
+import os
+import tensorflow as tf
+import tensorflow_compression as tfc
+import argparse
+from glob import glob
+
+##process the image into suitable dimension
+def load_img(path):
+    string = tf.io.read_file(path)
+    image = tf.image.decode_image(string, channels=3)
+    return image
+
+
+def load_model(args):
+    model = tf.keras.models.load_model(args.model_path,compile=False)
+    return model
+
+
+def compress(args):
+    model = load_model(args)
+
+
+    os.makedirs('outputs/binary', exist_ok=True)
+
+    if os.path.isdir(args.image_path):
+        pathes = glob(os.path.join(args.image_path, '*'))
+    else:
+        pathes = [args.image_path]
+
+    for path in pathes:
+        bitpath = "outputs/binary/{}.pth".format(os.path.basename(path).split('.')[0])
+
+        image = load_img(path)
+        compressed = model.compress(image)
+        packed = tfc.PackedTensors()
+        packed.pack(compressed)
+        with open(bitpath, "wb") as f:
+            f.write(packed.string)
+        num_pixels = tf.reduce_prod(tf.shape(image)[:-1])
+        bpp = len(packed.string) * 8 / num_pixels
+
+
+
+
+        print('=============================================================')
+        print(os.path.basename(path))
+
+        print('bitrate : {0:.4}bpp'.format(bpp))
+        print('=============================================================\n')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('model_path',type=str, default='final_model')
+    parser.add_argument('image_path',type=str, default='kodak/kodim20.png')
+
+    args = parser.parse_args()
+
+    compress(args)

decompress.py

@@ -0,0 +1,48 @@
+import os
+import tensorflow as tf
+import tensorflow_compression as tfc
+import argparse
+from glob import glob
+
+def load_model(args):
+    model = tf.keras.models.load_model(args.model_path,compile=False)
+    return model
+
+def decompress(model, args):
+    os.makedirs("outputs/reconstruction/", exist_ok=True)
+
+    if os.path.isdir(args.binary_path):
+        pathes = glob(os.path.join(args.binary_path, '*'))
+    else:
+        pathes = [args.binary_path]
+
+    for path in pathes:
+
+        print('========================================================================')
+        print('image', os.path.basename(path))
+
+        with open(path, "rb") as f:
+            packed = tfc.PackedTensors(f.read())
+        tensors = packed.unpack(dtypes)
+        x_hat = model.decompress(*tensors)
+
+
+        fakepath = "./outputs/reconstruction/{}.png".format(os.path.basename(path).split('.')[0])
+        string = tf.image.encode_png(x_hat)
+        tf.io.write_file(fakepath, string)
+
+
+        print('========================================================================\n')
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('model_path')
+    parser.add_argument('binary_path')
+
+
+    args = parser.parse_args()
+
+    model = load_model(args)
+    dtypes = [t.dtype for t in model.decompress.input_signature]
+
+    decompress(model, args)