From 5efe5b93da4e3d21260e3d7775efb6032d80404e Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Fri, 7 Oct 2022 15:32:39 +1000
Subject: [PATCH 01/13] Create folder and set up files to use in assignment

---
 recognition/45853757-VQVAE/README.MD  | 0
 recognition/45853757-VQVAE/dataset.py | 0
 recognition/45853757-VQVAE/modules.py | 0
 recognition/45853757-VQVAE/predict.py | 0
 recognition/45853757-VQVAE/train.py   | 0
 5 files changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 recognition/45853757-VQVAE/README.MD
 create mode 100644 recognition/45853757-VQVAE/dataset.py
 create mode 100644 recognition/45853757-VQVAE/modules.py
 create mode 100644 recognition/45853757-VQVAE/predict.py
 create mode 100644 recognition/45853757-VQVAE/train.py

diff --git a/recognition/45853757-VQVAE/README.MD b/recognition/45853757-VQVAE/README.MD
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/recognition/45853757-VQVAE/dataset.py b/recognition/45853757-VQVAE/dataset.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/recognition/45853757-VQVAE/modules.py b/recognition/45853757-VQVAE/modules.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/recognition/45853757-VQVAE/predict.py b/recognition/45853757-VQVAE/predict.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/recognition/45853757-VQVAE/train.py b/recognition/45853757-VQVAE/train.py
new file mode 100644
index 0000000000..e69de29bb2

From 435129e5b132fcda1736ccceb085d0a7c42a2c29 Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Fri, 21 Oct 2022 01:23:12 +1000
Subject: [PATCH 02/13] Added load_data() function

---
 recognition/45853757-VQVAE/dataset.py | 41 +++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/recognition/45853757-VQVAE/dataset.py b/recognition/45853757-VQVAE/dataset.py
index e69de29bb2..334c939cf5 100644
--- a/recognition/45853757-VQVAE/dataset.py
+++ b/recognition/45853757-VQVAE/dataset.py
@@ -0,0 +1,41 @@
+import numpy as np
+from PIL import Image
+import os
+import matplotlib.pyplot as plt
+
+def load_data():
+    """
+    Creates lists of the directory links for all of the images to be used in the model
+
+    Params: None
+
+    Returns: Three list objects containing sets of data for training, validation and testing.
+    """
+    # Initialise three empty lists for our data to be stored appropriately
+    training_data = []
+    validation_data = []
+    testing_data = []
+
+    # Create tuple pairs of the directories for the files with the images, and
+    # which list they should be sorted into
+    location_and_data_category = [("D:/keras_png_slices_data/keras_png_slices_train", \
+            training_data), ("D:/keras_png_slices_data/keras_png_slices_vaildate", \
+            validation_data), ("D:/keras_png_slices_data/keras_png_slices_test", \
+            testing_data)]
+
+    # Find and store the individual directories for each image in each file
+    for data_set in location_and_data_category:
+        for file_name in os.listdir(data_set[0]):
+            data_set[1].append(os.path.join(data_set[0], file_name))
+
+    return training_data, validation_data, testing_data
+
+
+# for file_path in os.listdir("D:/dog"):
+#     image_path = os.path.join("D:/dog", file_path)
+# img = Image.open(image_path)
+
+# img = Image.open("D:/dog/dog.png")
+# plt.imshow(img)
+# plt.show(block=True)
+

From cb387cd00e3f27e3406b37a253fd28e2ba4b3272 Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 15:52:56 +1000
Subject: [PATCH 03/13] Added preprocess_data() and updated load_data()

---
 recognition/45853757-VQVAE/dataset.py | 62 +++++++++++++++++----------
 1 file changed, 40 insertions(+), 22 deletions(-)

diff --git a/recognition/45853757-VQVAE/dataset.py b/recognition/45853757-VQVAE/dataset.py
index 334c939cf5..361d4772be 100644
--- a/recognition/45853757-VQVAE/dataset.py
+++ b/recognition/45853757-VQVAE/dataset.py
@@ -1,15 +1,42 @@
 import numpy as np
-from PIL import Image
+#from PIL import Image
+#import tensorflow as tf
+from tensorflow.keras.utils import load_img, img_to_array
 import os
 import matplotlib.pyplot as plt
 
+def preprocess_data(training_data, validation_data, testing_data):
+    """
+    Normalise each data set and find the variance of training data
+
+    Params: 
+
+    Returns: 
+    """
+    training_data = np.array(training_data)
+    training_data = training_data.astype('float16') / 255.
+    training_data = training_data.reshape((len(training_data), np.prod(training_data.shape[1:])))
+
+    validation_data = np.array(validation_data)
+    validation_data = validation_data.astype('float16') / 255.
+    validation_data = validation_data.reshape((len(validation_data), np.prod(validation_data.shape[1:])))
+
+    testing_data = np.array(testing_data)
+    testing_data = testing_data.astype('float16') / 255.
+    testing_data = testing_data.reshape((len(testing_data), np.prod(testing_data.shape[1:])))
+
+    variance = np.var(training_data / 255.)
+
+    return training_data, validation_data, testing_data, variance
+
+
 def load_data():
     """
-    Creates lists of the directory links for all of the images to be used in the model
+    Loads the data to be used for training, validating and testing the model.
 
     Params: None
 
-    Returns: Three list objects containing sets of data for training, validation and testing.
+    Returns: Three normalised data sets of images for training, validation and testing and the variance of the training dataset.
     """
     # Initialise three empty lists for our data to be stored appropriately
     training_data = []
@@ -18,24 +45,15 @@ def load_data():
 
     # Create tuple pairs of the directories for the files with the images, and
     # which list they should be sorted into
-    location_and_data_category = [("D:/keras_png_slices_data/keras_png_slices_train", \
-            training_data), ("D:/keras_png_slices_data/keras_png_slices_vaildate", \
-            validation_data), ("D:/keras_png_slices_data/keras_png_slices_test", \
-            testing_data)]
+    location_and_data_category = [["D:/keras_png_slices_data/keras_png_slices_train", \
+            training_data], ["D:/keras_png_slices_data/keras_png_slices_vaildate", \
+            validation_data], ["D:/keras_png_slices_data/keras_png_slices_test", \
+            testing_data]]
 
     # Find and store the individual directories for each image in each file
-    for data_set in location_and_data_category:
-        for file_name in os.listdir(data_set[0]):
-            data_set[1].append(os.path.join(data_set[0], file_name))
-
-    return training_data, validation_data, testing_data
-
-
-# for file_path in os.listdir("D:/dog"):
-#     image_path = os.path.join("D:/dog", file_path)
-# img = Image.open(image_path)
-
-# img = Image.open("D:/dog/dog.png")
-# plt.imshow(img)
-# plt.show(block=True)
-
+    for dataset in location_and_data_category:
+        for file_name in os.listdir(dataset[0]):
+            dataset[1].append(img_to_array(load_img(os.path.join(dataset[0], file_name), color_mode="grayscale")))
+        #dataset[1] = np.array(dataset[1]).squeeze()
+    
+    return preprocess_data(training_data, validation_data, testing_data)

From cb476cdfbfbdbe0793bf93ed62348ec8bb9c8914 Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 19:15:29 +1000
Subject: [PATCH 04/13] Add create_encoder and create_decoder functions,
 scaffolding for further functions.

---
 recognition/45853757-VQVAE/modules.py | 50 +++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/recognition/45853757-VQVAE/modules.py b/recognition/45853757-VQVAE/modules.py
index e69de29bb2..27f87344b0 100644
--- a/recognition/45853757-VQVAE/modules.py
+++ b/recognition/45853757-VQVAE/modules.py
@@ -0,0 +1,50 @@
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import layers, models
+
+import numpy as np
+
+length = 256*256
+batch_size = 32
+
+depth = 16
+kernel = 3
+
+
+def create_encoder(latent_dim=16):
+    """ Create a simple encoder """
+    encoder = layers.Sequential(name="encoder")
+    encoder.add(layers.Conv2D(depth, kernel, activation="relu", strides=2, padding="same", input_shape=(length,)))
+    encoder.add(layers.Conv2D(depth*2, kernel, activation="relu", strides=2, padding="same"))
+    encoder.add(layers.Conv2D(depth*4, kernel, activation="relu", strides=2, padding="same"))
+    encoder.add(layers.Conv2D(depth*8, kernel, activation="relu", strides=2, padding="same"))
+    encoder.add(layers.Conv2D(latent_dim, 1, padding="same"))
+    return encoder
+        
+def create_decoder():
+    """ Create a simple decoder """
+    decoder = layers.Sequential(name="decoder")
+    decoder.add(layers.Conv2D(depth*8, kernel, activation="relu", strides=2, padding="same"))
+    decoder.add(layers.Conv2D(depth*4, kernel, activation="relu", strides=2, padding="same"))
+    decoder.add(layers.Conv2D(depth*2, kernel, activation="relu", strides=2, padding="same"))
+    decoder.add(layers.Conv2D(depth, kernel, activation="relu", strides=2, padding="same"))
+    decoder.add(layers.Conv2D(1, kernel, padding="same"))
+    return decoder
+
+
+class VQLayer(layers.Layer):
+    pass
+    # Create the vector quantization layer for the model
+    
+
+class VQVAEModel(models.Sequential):
+    pass
+    # Use Sequential as base since i'm using sequential for encoder/decoder. 
+    # Source I'm referencing uses models.Model so adapt appropriately
+
+class PixelCNN:
+    pass
+    # Work out whether this needs to be implemented here or in a diff section/helper methods etc
+
+def do_training():
+    pass

From 1698391d8dcb86871e77d9ecf9c28432f653745d Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 20:53:22 +1000
Subject: [PATCH 05/13] Classes VQLayer and VQVAEModel roughly filled out

---
 recognition/45853757-VQVAE/modules.py | 95 ++++++++++++++++++++++++---
 1 file changed, 86 insertions(+), 9 deletions(-)

diff --git a/recognition/45853757-VQVAE/modules.py b/recognition/45853757-VQVAE/modules.py
index 27f87344b0..ab6b5c3965 100644
--- a/recognition/45853757-VQVAE/modules.py
+++ b/recognition/45853757-VQVAE/modules.py
@@ -4,13 +4,14 @@
 
 import numpy as np
 
-length = 256*256
+
 batch_size = 32
+epochs = 5
 
+length = 256*256
 depth = 16
 kernel = 3
 
-
 def create_encoder(latent_dim=16):
     """ Create a simple encoder """
     encoder = layers.Sequential(name="encoder")
@@ -33,18 +34,94 @@ def create_decoder():
 
 
 class VQLayer(layers.Layer):
-    pass
-    # Create the vector quantization layer for the model
-    
+    def __init__(self, n_embeddings, embedding_dim, beta=0.25, **kwargs):
+        super().__init__(**kwargs)
+        self.embedding_dim = embedding_dim
+        self.n_embeddings = n_embeddings
+        self.beta = beta
+
+        # Initialise embeddings
+        w_init = tf.random_uniform_initializer()
+        self.embeddings = tf.Variable(
+            initial_value=w_init(shape=(self.embedding_dim, self.n_embeddings), 
+            dtype="float32"), trainable=True, name="vqvae_embeddings"
+        )
+
+    def call(self, x):
+        # Calc then flatten inputs, not incl embedding dimension
+        input_shape = tf.shape(x)
+        flattened = tf.reshape(x, [-1, self.embedding_dim])
+
+        # Perform quantisation then reshape quantised values to orig shape
+        encoding_indices = self.get_code_indices(flattened)
+        encodings = tf.one_hot(encoding_indices, self.n_embeddings)
+        quantised = tf.matmul(encodings, self.embeddings, transponse_b=True)
+        quantised = tf.reshape(quantised, input_shape)
+        
+        # Vector quntisation loss is added to the layer
+        commitment_loss = tf.reduce_mean((tf.stop_gradient(quantised) - x) ** 2)
+        codebook_loss = tf.reduce_mean((quantised - tf.stop_gradient(x)) ** 2)
+        self.add_loss(self.beta * commitment_loss - codebook_loss)
+        
+        # Straight-through estimator 
+        return x + tf.stop_gradient(quantised - x)
+
+    def get_code_indices(self, flattened_inputs):
+        # Calc L2-normalised dist between inputs and codes
+        similarity = tf.matmul(flattened_inputs, self.embeddings)
+        distances = (
+            tf.reduce_sum(flattened_inputs ** 2, axis=1, keepdims=True) 
+            + tf.reduce_sum(self.embeddings ** 2, axis=0) - 2 * similarity
+        )
+
+        encoding_indices = tf.argmin(distances, axis=1)
+        return encoding_indices
+
 
 class VQVAEModel(models.Sequential):
-    pass
-    # Use Sequential as base since i'm using sequential for encoder/decoder. 
-    # Source I'm referencing uses models.Model so adapt appropriately
+    def __init__(self, variance, latent_dim, n_embeddings, **kwargs):
+        super(VQVAEModel, self).__init__(**kwargs)
+        self.variance = variance
+        self.latent_dim = latent_dim
+        self.n_embeddings = n_embeddings
+
+        # Build our model
+        self.add(VQLayer(n_embeddings, latent_dim, name="vector quantiser")) 
+        self.add(create_encoder(latent_dim))
+        self.add(create_decoder())
+
+        # Measure our losses
+        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
+        self.reconstruction_loss_tracker = keras.metrics.Mean(name="recontruction_loss")
+        self.vq_loss_tracker = keras.metrics.Mean(name="vq_loss")
+
+    @property
+    def metrics(self):
+        return [self.total_loss_tracker, self.reconstruction_loss_tracker, self.vq_loss_tracker]
+
+    def train_step(self, x):
+        with tf.GradientTape() as tape:
+            reconstructions = self.call(x)
+
+            reconstruction_loss = (tf.reduce_mean((x - reconstructions) ** 2) / self.variance)
+            total_loss = reconstruction_loss + sum(self.losses)
+
+        grads = tape.gradient(total_loss, self.trainable_variables)
+        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
+
+        self.total_loss_tracker.update_state(total_loss)
+        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
+        self.vq_loss_tracker.update_state(sum(self.losses))
+
+        return {
+            "loss": self.total_loss_tracker.result(),
+            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
+            "vqvae_loss": self.vq_loss_tracker.result()
+        }
+
 
 class PixelCNN:
     pass
-    # Work out whether this needs to be implemented here or in a diff section/helper methods etc
 
 def do_training():
     pass

From 8c88bd783fdab93ac66ab7b18f2e91675e8b102f Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 21:03:53 +1000
Subject: [PATCH 06/13] Removed function do_training as its functionality is
 implemented in train.py

---
 recognition/45853757-VQVAE/modules.py | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/recognition/45853757-VQVAE/modules.py b/recognition/45853757-VQVAE/modules.py
index ab6b5c3965..8f5088eb5e 100644
--- a/recognition/45853757-VQVAE/modules.py
+++ b/recognition/45853757-VQVAE/modules.py
@@ -1,13 +1,8 @@
 import tensorflow as tf
 from tensorflow import keras
 from tensorflow.keras import layers, models
-
 import numpy as np
 
-
-batch_size = 32
-epochs = 5
-
 length = 256*256
 depth = 16
 kernel = 3
@@ -118,10 +113,3 @@ def train_step(self, x):
             "reconstruction_loss": self.reconstruction_loss_tracker.result(),
             "vqvae_loss": self.vq_loss_tracker.result()
         }
-
-
-class PixelCNN:
-    pass
-
-def do_training():
-    pass

From 9c683b13bfd5934c552de2df7ea04e5959670d3b Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 21:42:32 +1000
Subject: [PATCH 07/13] Updated preprocess_data(), cleaned docstrings/imports

---
 recognition/45853757-VQVAE/dataset.py | 23 ++++++-----------------
 1 file changed, 6 insertions(+), 17 deletions(-)

diff --git a/recognition/45853757-VQVAE/dataset.py b/recognition/45853757-VQVAE/dataset.py
index 361d4772be..15f315d61d 100644
--- a/recognition/45853757-VQVAE/dataset.py
+++ b/recognition/45853757-VQVAE/dataset.py
@@ -1,29 +1,19 @@
+import os
 import numpy as np
-#from PIL import Image
-#import tensorflow as tf
 from tensorflow.keras.utils import load_img, img_to_array
-import os
-import matplotlib.pyplot as plt
 
-def preprocess_data(training_data, validation_data, testing_data):
-    """
-    Normalise each data set and find the variance of training data
 
-    Params: 
+def preprocess_data(training_data, validation_data, testing_data):
+    " Normalises each data set and finds the variance of training data "
 
-    Returns: 
-    """
     training_data = np.array(training_data)
     training_data = training_data.astype('float16') / 255.
-    training_data = training_data.reshape((len(training_data), np.prod(training_data.shape[1:])))
 
     validation_data = np.array(validation_data)
     validation_data = validation_data.astype('float16') / 255.
-    validation_data = validation_data.reshape((len(validation_data), np.prod(validation_data.shape[1:])))
-
+   
     testing_data = np.array(testing_data)
     testing_data = testing_data.astype('float16') / 255.
-    testing_data = testing_data.reshape((len(testing_data), np.prod(testing_data.shape[1:])))
 
     variance = np.var(training_data / 255.)
 
@@ -43,17 +33,16 @@ def load_data():
     validation_data = []
     testing_data = []
 
-    # Create tuple pairs of the directories for the files with the images, and
+    # Create list pairs of the directories for the files with the images, and
     # which list they should be sorted into
     location_and_data_category = [["D:/keras_png_slices_data/keras_png_slices_train", \
             training_data], ["D:/keras_png_slices_data/keras_png_slices_vaildate", \
             validation_data], ["D:/keras_png_slices_data/keras_png_slices_test", \
             testing_data]]
 
-    # Find and store the individual directories for each image in each file
+    # Find and store each image in each file into the correct list
     for dataset in location_and_data_category:
         for file_name in os.listdir(dataset[0]):
             dataset[1].append(img_to_array(load_img(os.path.join(dataset[0], file_name), color_mode="grayscale")))
-        #dataset[1] = np.array(dataset[1]).squeeze()
     
     return preprocess_data(training_data, validation_data, testing_data)

From 7537f534a531690139f8aecd3f8938c8d31d7b68 Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 22:55:11 +1000
Subject: [PATCH 08/13] Bugs: Fixed issues with importing Sequential

---
 recognition/45853757-VQVAE/modules.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/recognition/45853757-VQVAE/modules.py b/recognition/45853757-VQVAE/modules.py
index 8f5088eb5e..f451692efc 100644
--- a/recognition/45853757-VQVAE/modules.py
+++ b/recognition/45853757-VQVAE/modules.py
@@ -3,14 +3,14 @@
 from tensorflow.keras import layers, models
 import numpy as np
 
-length = 256*256
+length = 256
 depth = 16
 kernel = 3
 
 def create_encoder(latent_dim=16):
     """ Create a simple encoder """
-    encoder = layers.Sequential(name="encoder")
-    encoder.add(layers.Conv2D(depth, kernel, activation="relu", strides=2, padding="same", input_shape=(length,)))
+    encoder = tf.keras.Sequential(name="encoder")
+    encoder.add(layers.Conv2D(depth, kernel, activation="relu", strides=2, padding="same", input_shape=(length, length, 1)))
     encoder.add(layers.Conv2D(depth*2, kernel, activation="relu", strides=2, padding="same"))
     encoder.add(layers.Conv2D(depth*4, kernel, activation="relu", strides=2, padding="same"))
     encoder.add(layers.Conv2D(depth*8, kernel, activation="relu", strides=2, padding="same"))
@@ -19,7 +19,7 @@ def create_encoder(latent_dim=16):
         
 def create_decoder():
     """ Create a simple decoder """
-    decoder = layers.Sequential(name="decoder")
+    decoder = tf.keras.Sequential(name="decoder")
     decoder.add(layers.Conv2D(depth*8, kernel, activation="relu", strides=2, padding="same"))
     decoder.add(layers.Conv2D(depth*4, kernel, activation="relu", strides=2, padding="same"))
     decoder.add(layers.Conv2D(depth*2, kernel, activation="relu", strides=2, padding="same"))
@@ -73,7 +73,7 @@ def get_code_indices(self, flattened_inputs):
         return encoding_indices
 
 
-class VQVAEModel(models.Sequential):
+class VQVAEModel(tf.keras.Sequential):
     def __init__(self, variance, latent_dim, n_embeddings, **kwargs):
         super(VQVAEModel, self).__init__(**kwargs)
         self.variance = variance

From b6a98e09b8565ec34c3d55885b10f09813ffcc86 Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 23:28:48 +1000
Subject: [PATCH 09/13] Docstrings updated

---
 recognition/45853757-VQVAE/dataset.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/recognition/45853757-VQVAE/dataset.py b/recognition/45853757-VQVAE/dataset.py
index 15f315d61d..47f220926e 100644
--- a/recognition/45853757-VQVAE/dataset.py
+++ b/recognition/45853757-VQVAE/dataset.py
@@ -4,11 +4,23 @@
 
 
 def preprocess_data(training_data, validation_data, testing_data):
-    " Normalises each data set and finds the variance of training data "
-
+    """ 
+    Normalises each data set and finds the variance of training data
+    
+    Parameters:
+        training_data (list): list of arrays representing the images to train the model on  
+        validation_data (list): list of arrays representing the images to validate the model on  
+        testing_data (list): list of arrays representing the images to test the model on  
+
+    Returns:
+        training_data (list): normalised list of arrays
+        validation_data (list): normalised list of arrays  
+        testing_data (list): normalised list of arrays
+        variance (float): variance of the training data
+    """
     training_data = np.array(training_data)
     training_data = training_data.astype('float16') / 255.
-
+  
     validation_data = np.array(validation_data)
     validation_data = validation_data.astype('float16') / 255.
    
@@ -26,7 +38,9 @@ def load_data():
 
     Params: None
 
-    Returns: Three normalised data sets of images for training, validation and testing and the variance of the training dataset.
+    Returns: 
+        Three normalised data sets of images for training, 
+        validation and testing and the variance of the training dataset.
     """
     # Initialise three empty lists for our data to be stored appropriately
     training_data = []

From c0cd4d75c4d5c26b86e0058311339370e3222605 Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 23:37:43 +1000
Subject: [PATCH 10/13] Added train_vqvae()

---
 recognition/45853757-VQVAE/train.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/recognition/45853757-VQVAE/train.py b/recognition/45853757-VQVAE/train.py
index e69de29bb2..2764910240 100644
--- a/recognition/45853757-VQVAE/train.py
+++ b/recognition/45853757-VQVAE/train.py
@@ -0,0 +1,17 @@
+import tensorflow as tf
+from modules import *
+from dataset import *
+
+batch_size = 64
+epochs = 100
+
+def train_vqvae():
+    # Load our data
+    training_data, validation_data, testing_data, data_variance = load_data()
+
+    # Construct and train our model
+    vqvae_model = VQVAEModel(variance=data_variance, latent_dim=16, n_embeddings=128)
+    vqvae_model.compile(optimizer=keras.optimizers.Adam())
+    print(vqvae_model.summary())
+
+    return vqvae_model.fit(training_data, training_data, validation_data=(validation_data, validation_data), epochs=epochs, batch_size=batch_size)

From 65b93dc4a1dc718ee05f1de046596601c9f56f33 Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 23:48:25 +1000
Subject: [PATCH 11/13] Added calculate_ssim(), compare_predicted(), and code
 to build the model

---
 recognition/45853757-VQVAE/predict.py | 33 +++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/recognition/45853757-VQVAE/predict.py b/recognition/45853757-VQVAE/predict.py
index e69de29bb2..c884643213 100644
--- a/recognition/45853757-VQVAE/predict.py
+++ b/recognition/45853757-VQVAE/predict.py
@@ -0,0 +1,33 @@
+from dataset import *
+from modules import *
+from train import *
+import tensorflow as tf
+import matplotlib.pyplot as plt
+
+def calculate_ssim(original_data, predicted_data):
+    """ Calculates and the average of the SSIM of all images in the two sets of data as a percentage. """
+    ssim = tf.image.ssim(original_data, predicted_data, max_val=1)
+    print("SSIM of data sets:", ssim)
+
+def compare_predicted(original_data, predicted_data, index):
+    """ Plots the original and predicted image at the given index """
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 2, 1)
+    imgplot = plt.imshow(original_data[index])
+    ax.set_title("Original Image")
+
+    ax = fig.add_subplot(1, 2, 2)
+    imgplot = plt.imshow(predicted_data[index])
+    ax.set_title("Reconstructed Image")
+    fig.show()
+
+def plot_loss(model):
+    return None
+
+training_data, validation_data, testing_data, data_variance = load_data()
+model = train_vqvae()
+
+predictions = model.predict(testing_data)
+
+calculate_ssim(testing_data, predictions)
+compare_predicted(testing_data, predictions, 8)
\ No newline at end of file

From b1ec37d2ce8e8cf6d606e916746693c1f5f372c3 Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 23:51:55 +1000
Subject: [PATCH 12/13] Updated docstrings

---
 recognition/45853757-VQVAE/modules.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/recognition/45853757-VQVAE/modules.py b/recognition/45853757-VQVAE/modules.py
index f451692efc..b8b2b16206 100644
--- a/recognition/45853757-VQVAE/modules.py
+++ b/recognition/45853757-VQVAE/modules.py
@@ -8,7 +8,7 @@
 kernel = 3
 
 def create_encoder(latent_dim=16):
-    """ Create a simple encoder """
+    """ Create a simple encoder sequential layer """
     encoder = tf.keras.Sequential(name="encoder")
     encoder.add(layers.Conv2D(depth, kernel, activation="relu", strides=2, padding="same", input_shape=(length, length, 1)))
     encoder.add(layers.Conv2D(depth*2, kernel, activation="relu", strides=2, padding="same"))
@@ -18,7 +18,7 @@ def create_encoder(latent_dim=16):
     return encoder
         
 def create_decoder():
-    """ Create a simple decoder """
+    """ Create a simple decoder sequential layer """
     decoder = tf.keras.Sequential(name="decoder")
     decoder.add(layers.Conv2D(depth*8, kernel, activation="relu", strides=2, padding="same"))
     decoder.add(layers.Conv2D(depth*4, kernel, activation="relu", strides=2, padding="same"))
@@ -69,6 +69,7 @@ def get_code_indices(self, flattened_inputs):
             + tf.reduce_sum(self.embeddings ** 2, axis=0) - 2 * similarity
         )
 
+        # Derive for min distances
         encoding_indices = tf.argmin(distances, axis=1)
         return encoding_indices
 
@@ -95,15 +96,17 @@ def metrics(self):
         return [self.total_loss_tracker, self.reconstruction_loss_tracker, self.vq_loss_tracker]
 
     def train_step(self, x):
+        # Calculates the losses from the VQ-VAE
         with tf.GradientTape() as tape:
             reconstructions = self.call(x)
-
             reconstruction_loss = (tf.reduce_mean((x - reconstructions) ** 2) / self.variance)
             total_loss = reconstruction_loss + sum(self.losses)
 
+        # Backpropagates our values
         grads = tape.gradient(total_loss, self.trainable_variables)
         self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
 
+        # Implement loss tracking
         self.total_loss_tracker.update_state(total_loss)
         self.reconstruction_loss_tracker.update_state(reconstruction_loss)
         self.vq_loss_tracker.update_state(sum(self.losses))

From b3298bba385d2525dcb20492e514d3422b02b213 Mon Sep 17 00:00:00 2001
From: Danielle Soilleux <d.soilleux@uqconnect.edu.au>
Date: Sat, 22 Oct 2022 23:55:08 +1000
Subject: [PATCH 13/13] Added text to README

---
 recognition/45853757-VQVAE/README.MD | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/recognition/45853757-VQVAE/README.MD b/recognition/45853757-VQVAE/README.MD
index e69de29bb2..c597b0c3c1 100644
--- a/recognition/45853757-VQVAE/README.MD
+++ b/recognition/45853757-VQVAE/README.MD
@@ -0,0 +1,8 @@
+# Implementing a VQ-VAE for the OASIS dataset
+VQ-VAEs are a variety of variational autoencoders (hence VAE), which take on the ideas of vector quantisation (VQ) to improve their effectiveness. 
+
+My VQVAE model was constructed based on the implementation from Keras at: https://keras.io/examples/generative/vq_vae/
+
+The OASIS dataset was downloaded through blackboard and preprocessing it included normalising the images. 
+
+I did not get very conclusive results as my GPU cannot handle this many data samples and I could not finid an adequate solution.