Set pretrained layer back to trainable

- Was actually quite easy

eyeballer/model.py

@@ -1,6 +1,7 @@
 import os
 import random
 import sys
+import progressbar
 
 # Prevent Tkinter Dependency
 import matplotlib
@@ -10,16 +11,9 @@ import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import Augmentor
+import tensorflow as tf
 
-from keras import regularizers
-from keras.applications import MobileNet
-from keras.applications.mobilenet import preprocess_input
-from keras.callbacks import ModelCheckpoint
-from keras.layers import Dropout, Dense
-from keras.layers import GlobalAveragePooling2D
-from keras.models import Model
-from keras.optimizers import Adam
-from keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
 from sklearn.metrics import classification_report, accuracy_score, hamming_loss
 from eyeballer.augmentation import EyeballerAugmentation
 
@@ -45,23 +39,24 @@ class EyeballModel:
         weights_file -- A filename for where to load the model's weights.
         seed -- PRNG seed, useful for repeating a previous run and using the same data. Training/Validation split is determined randomly.
         """
-        # Build the model
-        base_model = MobileNet(weights='imagenet', include_top=False, input_shape=self.input_shape)
-        x = base_model.output
-        x = GlobalAveragePooling2D()(x)
-        x = Dense(256, activation="relu", name="HiddenLayer1")(x)
-        x = Dropout(0.5)(x)
-        x = Dense(128, activation="relu", name="HiddenLayer2", kernel_regularizer=regularizers.l2(0.01))(x)
-        x = Dropout(0.2)(x)
-        output_layer = Dense(len(DATA_LABELS), activation="sigmoid", name="OutputLayer")(x)
-        self.model = Model(inputs=base_model.input, outputs=output_layer)
+        # # Build the model
+        self.model = tf.keras.Sequential()
+        pretrained_layer = tf.keras.applications.mobilenet_v2.MobileNetV2(weights='imagenet', include_top=False, input_shape=self.input_shape)
+        self.model.add(pretrained_layer)
+        self.model.add(tf.keras.layers.GlobalAveragePooling2D())
+        self.model.add(tf.keras.layers.Dense(256, activation="relu"))
+        self.model.add(tf.keras.layers.Dropout(0.5))
+        self.model.add(tf.keras.layers.Dense(128, activation="relu"))
+        self.model.add(tf.keras.layers.Dropout(0.2))
+        self.model.add(tf.keras.layers.Dense(len(DATA_LABELS), activation="sigmoid"))
 
-        # for layer in base_model.layers:
+        self.model.compile(optimizer=tf.keras.optimizers.Adam(0.0005),
+                           loss="binary_crossentropy",
+                           metrics=["accuracy"])
+
+        # for layer in pretrained_layer.layers:
         #     layer.trainable = False
 
-        adam = Adam(lr=0.0005)
-        self.model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
-
         if print_summary:
             print(self.model.summary())
 
@@ -117,7 +112,7 @@ class EyeballModel:
         """
         print("Training with seed: " + str(self.seed))
 
-        data_generator = ImageDataGenerator(
+        data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
             preprocessing_function=self.preprocess_training_function,
             validation_split=0.2,
             samplewise_center=True)
@@ -146,8 +141,8 @@ class EyeballModel:
             class_mode="other")
 
         # Model checkpoint - Saves model weights when validation accuracy improves
-        callbacks = [ModelCheckpoint(self.checkpoint_file, monitor='val_loss', verbose=1, save_best_only=True,
-                                     save_weights_only=True, mode='min')]
+        callbacks = [tf.keras.callbacks.ModelCheckpoint(self.checkpoint_file, monitor='val_loss', verbose=1, save_best_only=True,
+                     save_weights_only=True, mode='min')]
 
         history = self.model.fit_generator(
             training_generator,
@@ -201,11 +196,10 @@ class EyeballModel:
         return result
 
     def predict(self, path, threshold=0.5):
-        """Predict the labels for a single file
+        """Predict the labels for a single file or directory of files
 
         Keyword arguments:
-        path -- The path to the file that we'll be evaluating.
-            NOTE: This actually has to be a directory containing a directory with our files. Because Keras is dumb. Don't ask.
+        path -- The path to the file(s) that we'll be evaluating.
         """
         # Is this a single file, or a directory?
         screenshots = []
@@ -217,29 +211,26 @@ class EyeballModel:
         else:
             raise FileNotFoundError
 
-        # Prepare the data
-        data_generator = ImageDataGenerator(
-            preprocessing_function=preprocess_input)
-        prediction_generator = data_generator.flow_from_directory(
-            directory=path,
-            target_size=(self.image_width, self.image_height),
-            shuffle=False,
-            class_mode=None)
-
-        predictions = self.model.predict_generator(
-            prediction_generator,
-            verbose=1,
-            steps=len(prediction_generator))
-
-        predictions = predictions > threshold
         results = []
-        for filename, prediction in zip(prediction_generator.filenames, predictions):
+        for screenshot in progressbar.progressbar(screenshots):
+            # Load the image into memory
+            img = None
+            try:
+                img = tf.keras.preprocessing.image.load_img(screenshot, target_size=(self.image_width, self.image_height))
+                img = tf.keras.preprocessing.image.img_to_array(img)
+                img = np.expand_dims(img, axis=0)
+                img = preprocess_input(img)
+            except IsADirectoryError:
+                print("\nWARN: Skipping directory: ", screenshot)
+                continue
+
+            prediction = self.model.predict(img, batch_size=1)
             result = dict()
-            result["filename"] = path + filename
-            result["custom404"] = prediction[0]
-            result["login"] = prediction[1]
-            result["homepage"] = prediction[2]
-            result["oldlooking"] = prediction[3]
+            result["filename"] = screenshot
+            result["custom404"] = prediction[0][0]
+            result["login"] = prediction[0][1]
+            result["homepage"] = prediction[0][2]
+            result["oldlooking"] = prediction[0][3]
             results.append(result)
         return results
 
@@ -250,7 +241,7 @@ class EyeballModel:
         threshold -- Value between 0->1. The cutoff where the numerical prediction becomes boolean. Default: 0.5
         """
         # Prepare the data
-        data_generator = ImageDataGenerator(
+        data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
             preprocessing_function=preprocess_input)
         evaluation_generator = data_generator.flow_from_dataframe(
             self.evaluation_labels,
@@ -265,7 +256,7 @@ class EyeballModel:
         if not self.random_seed:
             print("Using validation set...")
             # Data augmentation
-            data_generator = ImageDataGenerator(
+            data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
                 preprocessing_function=self.preprocess_training_function,
                 samplewise_center=True,
                 validation_split=0.2)

requirements-gpu.txt

@@ -5,5 +5,6 @@ numpy
 pandas
 pillow
 sklearn
-tensorflow-gpu
-keras
+tensorflow-gpu==2.0.0-beta0
+jinja2
+progressbar2

requirements.txt

@@ -5,6 +5,6 @@ numpy
 pandas
 pillow
 sklearn
-tensorflow
-keras
+tensorflow==2.0.0-beta0
 jinja2
+progressbar2