Simplified models, config and plot

Author: Light4Code

README.md

@@ -27,8 +27,8 @@ python src/export.py config/sample.json --checkpoint_path PATH_TO_CHECKPOINT --o
 ## Models
 ### Anomaly detection
 Anomaly detection will use a autoencoder approche, the prediction substratced from the original image should show the anomaly.
-- Fast: Simple 1D prediction
-- Advanced: Simple 2D prediction
+- Deep Autoencoder: Simple 1D autoencoder
+- Convolutional Autoencoder
 ### Segmentation
 #### Unet
 - [Vanilla Unet](https://arxiv.org/pdf/1505.04597.pdf) (original paper)

config/sample.json

@@ -1,6 +1,6 @@
 {
   "input_shape": [512, 512, 3],
-  "model": "fast",
+  "model": "deep_autoencoder",
   "train": {
     "batch_size": 40,
     "epochs": 100,
@@ -20,7 +20,7 @@
       "height_shift_range": 0.3,
       "rotation_range": 10
     },
-    "fast_model": {
+    "deep_autoencoder_model": {
       "translator_layer_size": 200,
       "middle_layer_size": 32
     }

src/image_generator.py

@@ -96,6 +96,7 @@ def main():
         print("Finished round (" + str(r+1) + "/" + str(rounds) + ")")
 
 def generate_images(config, image_util, output_path, prefix, class_name, images, masks):
+    prefix = str(prefix) + "_" + class_name
     images_output_path = output_path + "/" + class_name
     masks_output_path = images_output_path + "/masks"
 

src/models/anomaly_detection/__init__.py

@@ -1,2 +1,2 @@
-from .advanced_model import AdvancedModel
-from .fast_model import FastModel
+from .convolutional_autoencoder_model import ConvolutionalAutoencoderModel
+from .deep_autoencoder_model import DeepAutoencoderModel

src/models/anomaly_detection/advanced_model.py src/models/anomaly_detection/convolutional_autoencoder_model.py

@@ -1,34 +1,43 @@
+import matplotlib.pyplot as plt
 import numpy as np
+import numpy.ma as ma
 import tensorflow as tf
 from tensorflow.keras import Input
-from tensorflow.keras.layers import (
-    Activation,
-    BatchNormalization,
-    Conv2D,
-    Conv2DTranspose,
-    Dense,
-    Flatten,
-    LeakyReLU,
-    Reshape,
-)
+from tensorflow.keras.layers import (Activation, BatchNormalization, Conv2D,
+                                     Conv2DTranspose, Dense, Flatten,
+                                     LeakyReLU, Reshape)
 from tensorflow.keras.models import Model
 from tensorflow.keras.optimizers import Adam
+
 from models import BaseModel
+from utils.plots import *
 
 
-class AdvancedModel(BaseModel):
+class ConvolutionalAutoencoderModel(BaseModel):
     def __init__(self, config):
         super().__init__(config)
 
     def create_optimizer(self, optimzer="adam"):
         super().create_optimizer(optimzer)
 
-    def create_model(self, filters=(32, 64), latent_dim=16):
+    def compile(self, loss="mse"):
+        self.model.compile(loss=loss, optimizer=self.optimizer, metrics=["accuracy"])
+
+    def create_model(self):
+        filters = (32, 64)
+        kernel_size = (3,3)
+        try:
+            model_config = self.config.train.raw["convolutional_autoencoder_model"]
+            latent_dim = model_config["latent_dim"]
+        except:
+            latent_dim = 16
+            
+
         input_shape = self.config.input_shape
         inputs = Input(shape=input_shape, name=self.input_name)
         x = inputs
         for f in filters:
-            x = Conv2D(filters=f, kernel_size=(3, 3), strides=2, padding="same")(x)
+            x = Conv2D(filters=f, kernel_size=kernel_size, strides=2, padding="same")(x)
             x = LeakyReLU(alpha=0.2)(x)
             x = BatchNormalization(axis=input_shape[2])(x)
         volume_size = tf.keras.backend.int_shape(x)
@@ -41,11 +50,11 @@ def create_model(self, filters=(32, 64), latent_dim=16):
         x = Reshape((volume_size[1], volume_size[2], volume_size[3]))(x)
         for f in filters[::-1]:
             x = Conv2DTranspose(
-                filters=f, kernel_size=(3, 3), strides=2, padding="same"
+                filters=f, kernel_size=kernel_size, strides=2, padding="same"
             )(x)
             x = LeakyReLU(alpha=0.2)(x)
             x = BatchNormalization(axis=input_shape[2])(x)
-        x = Conv2DTranspose(filters=input_shape[2], kernel_size=(3, 3), padding="same")(
+        x = Conv2DTranspose(filters=input_shape[2], kernel_size=kernel_size, padding="same")(
             x
         )
         outputs = Activation("sigmoid", name=self.output_name)(x)  # Decoded
@@ -54,9 +63,5 @@ def create_model(self, filters=(32, 64), latent_dim=16):
         self.model = Model(inputs, decoder(encoder(inputs)), name="autoencoder")
         return self.model
 
-    def overwrite_optimizer(self, optimizer, optimizer_name):
-        self.optimzer = optimizer
-        self.optimizer_name = optimizer_name
-
-    def compile(self, loss="mse"):
-        self.model.compile(loss=loss, optimizer=self.optimizer, metrics=["accuracy"])
+    def plot_predictions(self, test_images):
+        plot_difference(self.config, self.predictions, test_images)

src/models/anomaly_detection/deep_autoencoder_model.py

@@ -0,0 +1,53 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import numpy.ma as ma
+import tensorflow as tf
+from tensorflow.keras.layers import BatchNormalization, Dense, Flatten, Reshape
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
+
+from models import BaseModel
+from utils.plots import *
+
+
+class DeepAutoencoderModel(BaseModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+    def create_optimizer(self, optimzer="adam"):
+        super().create_optimizer(optimzer)
+
+    def create_model(self):
+        input_shape = self.config.input_shape
+
+        try:
+            model_config = self.config.train.raw["deep_autoencoder_model"]
+            translator_layer_size = model_config["translator_layer_size"]
+            middle_layer_size = model_config["middle_layer_size"]
+        except:
+            translator_layer_size = 100
+            middle_layer_size = 16
+
+        sub_layer_size = int(translator_layer_size / 2)
+        input_dim = input_shape[0] * input_shape[1] * input_shape[2]
+        input = tf.keras.Input(input_shape, name=self.input_name)
+        x = input
+        x = Flatten()(x)
+        x = BatchNormalization()(x)
+        x = Dense(translator_layer_size, activation="relu", name="encoder")(x)
+        x = Dense(sub_layer_size, activation="relu")(x)
+        x = Dense(middle_layer_size, activation="relu")(x)
+        x = Dense(sub_layer_size, activation="relu")(x)
+        x = BatchNormalization()(x)
+        x = Dense(translator_layer_size, activation="relu", name="decoder")(x)
+        x = Dense(input_dim, activation="sigmoid", name="reconstructor")(x)
+        x = Reshape(input_shape, name=self.output_name)(x)
+
+        self.model = Model(input, x)
+        return self.model
+
+    def compile(self, loss="mean_squared_error"):
+        self.model.compile(loss=loss, optimizer=self.optimizer, metrics=["accuracy"])
+
+    def plot_predictions(self, test_images):
+        plot_difference(self.config, self.predictions, test_images)

src/models/anomaly_detection/fast_model.py

@@ -10,6 +10,7 @@
 from tensorflow.keras.preprocessing.image import ImageDataGenerator
 
 from utils import *
+from utils.plots import *
 
 
 class BaseModel:
@@ -84,7 +85,7 @@ def load_weights(self):
 
     def train(self):
         if not self.train_datagen == None:
-            self.model.fit(
+            self.history = self.model.fit(
                 self.train_datagen,
                 epochs=self.config.train.epochs,
                 steps_per_epoch=len(self.train_images) / self.config.train.batch_size,
@@ -95,7 +96,7 @@ def train(self):
             )
 
         else:
-            self.model.fit(
+            self.history = self.model.fit(
                 self.train_images,
                 self.y_train,
                 batch_size=self.config.train.batch_size,
@@ -114,36 +115,7 @@ def predict(self, test_images):
             )
 
     def plot_predictions(self, test_images):
-        pred_count = len(self.predictions)
-        plt_shape = (self.config.input_shape[0], self.config.input_shape[1])
-        plt_cmap = "gray"
-        if self.config.input_shape[2] > 1:
-            plt_shape = (
-                self.config.input_shape[0],
-                self.config.input_shape[1],
-                self.config.input_shape[2],
-            )
-        index = 1
-        plt_index = 0
-        for test_image in test_images:
-            original_image = test_image.reshape(plt_shape)
-            pred_image = self.predictions[plt_index].reshape(plt_shape)
-            mask = ma.masked_where(pred_image < self.config.test_threshold, pred_image)
-            plt.subplot(pred_count, 3, index)
-            plt.title("Original")
-            plt.imshow(original_image, interpolation="none", cmap=plt_cmap)
-            index += 1
-            plt.subplot(pred_count, 3, index)
-            plt.title("Prediction")
-            plt.imshow(pred_image, interpolation="none", cmap=plt_cmap)
-            index += 1
-            plt.subplot(pred_count, 3, index)
-            plt.title("Overlay")
-            plt.imshow(original_image, interpolation="none", cmap=plt_cmap)
-            plt.imshow(mask, cmap="jet", interpolation="none", alpha=0.7)
-            index += 1
-            plt_index += 1
-        plt.show()
+        utils.plot_prediction(config, self.predictions, test_images)
 
     def prepare_training(self):
         self.train_images = None

src/models/base_model.py

@@ -82,10 +82,18 @@ def main():
         metavar="path",
         help="Overwrites the path to the saved checkpoint containing the model weights",
     )
+    parser.add_argument(
+        "--plot_history",
+        dest="plot_history",
+        metavar="boolean (default: false)",
+        type=bool,
+        help="Plots the model training history",
+    )
 
     args = parser.parse_args()
     config_path = args.config
     config = Config(config_path)
+    plot_history = False
 
     # Overwrite config
     if args.train_files_path:
@@ -109,6 +117,9 @@ def main():
     if args.checkpoint_path:
         config.train.checkpoint_path = args.checkpoint_path
 
+    if args.plot_history:
+        plot_history = True
+
     # Set seed to get reproducable experiments
     seed_value = 33
     os.environ["PYTHONHASHSEED"] = str(seed_value)
@@ -128,7 +139,26 @@ def main():
 
     # ToDo: Train model
     model.train()
+    history = model.history
+    epochs = len(history.epoch) + model.initial_epoch
+    model.model.save_weights(config.train.checkpoints_path + "/model-{0:04d}.ckpts".format(epochs))
 
+    if plot_history:
+        plt.subplot(121)
+        plt.plot(history.history['accuracy'])
+        plt.plot(history.history['val_accuracy'])
+        plt.title('Model accuracy')
+        plt.ylabel('Accuracy')
+        plt.xlabel('Epoch')
+        plt.legend(['Train', 'Test'], loc='upper left')
+        plt.subplot(122)
+        plt.plot(history.history['loss'])
+        plt.plot(history.history['val_loss'])
+        plt.title('Model loss')
+        plt.ylabel('Loss')
+        plt.xlabel('Epoch')
+        plt.legend(['Train', 'Test'], loc='upper left')
+        plt.show()
 
 if __name__ == "__main__":
     main()

src/train.py

@@ -1,3 +1,4 @@
 from .config import Config
 from .image_util import ImageUtil
 from .model_creater import *
+from .plots import *