Fail more loudly when model file has an error

FileNotFound and invalid models will now throw errors

.gitignore

@@ -113,3 +113,4 @@ venv.bak/
 *.model
 *.csv
 *.txt
+!requirements.txt

eyeballer/model.py

@@ -1,7 +1,12 @@
 import os
 import random
+import sys
 
+# Prevent Tkinter Dependency
+import matplotlib
+matplotlib.use('agg')  # noqa: E402
 import matplotlib.pyplot as plt
+
 import numpy as np
 import pandas as pd
 import Augmentor
@@ -18,10 +23,23 @@ from keras.preprocessing import image
 from keras.preprocessing.image import ImageDataGenerator
 from sklearn.metrics import classification_report, accuracy_score, hamming_loss
 from augmentation import EyeballerAugmentation
+from keras.callbacks import Callback
 
 DATA_LABELS = ["custom404", "login", "homepage", "oldlooking"]
 
 
+class PredictionManager(Callback):
+    def __init__(self):
+        super(PredictionManager, self).__init__()
+        self.predictions = []
+
+    def getPredictions(self):
+        return np.array(self.predictions)
+
+    def on_predict_batch_end(self, batch, logs=None):
+        self.predictions.append(logs['outputs'][0].tolist()[0])
+
+
 class EyeballModel:
     """The primary model class of Eyeballer.
 
@@ -33,7 +51,7 @@ class EyeballModel:
     image_width, image_height = 224, 224
     input_shape = (image_width, image_height, 3)
 
-    def __init__(self, print_summary=False, weights_file="weights.h5", seed=None):
+    def __init__(self, weights_file, print_summary=False, seed=None):
         """Constructor for model class.
 
         Keyword arguments:
@@ -78,10 +96,16 @@ class EyeballModel:
         self.training_labels = self.training_labels.sample(frac=1)
 
         if weights_file is not None and os.path.isfile(weights_file):
-            self.model.load_weights(weights_file)
+            try:
+                self.model.load_weights(weights_file)
+            except OSError as e:
+                print(f"ERROR: Unable to open weights file '{weights_file}'")
+                sys.exit(-1)
             print("Loaded model from file.")
         else:
-            print("No model loaded from file")
+            if weights_file is not None:
+                raise FileNotFoundError
+            print("WARN: No model loaded from file. Generating random model")
 
         # Data augmentation
         augmentor = Augmentor.Pipeline()
@@ -247,7 +271,22 @@ class EyeballModel:
         else:
             print("Using evaluation set...")
 
-        predictions = self.model.predict_generator(evaluation_generator, verbose=1, steps=len(evaluation_generator))
+        prediction_manager = PredictionManager()
+
+        try:
+            self.model.predict_generator(
+                evaluation_generator,
+                callbacks=[prediction_manager],
+                verbose=1,
+                steps=len(evaluation_generator))
+        except KeyboardInterrupt:
+            print('\r\nJob interrupted (Ctrl+C): Outputting completed results')
+
+        predictions = prediction_manager.getPredictions()
+
+        # Truncates evaluation generator if test was interrupted
+        evaluation_generator._data = evaluation_generator.data[:predictions.shape[0]]
+
         self._save_prediction_histograms(predictions)
         predictions = predictions > threshold
         ground_truth = evaluation_generator.data

requirements.txt

@@ -0,0 +1,9 @@
+Augmentor
+click
+matplotlib
+numpy
+pandas
+pillow
+sklearn
+tensorflow
+git+https://github.com/keras-team/keras/

setup.py

@@ -1,22 +0,0 @@
-from setuptools import find_packages, setup
-
-setup(
-    name='eyeballer',
-    version='0.0.1',
-    packages=find_packages(exclude=["tests"]),
-    url='https://github.com/BishopFox/eyeballer',
-    license='GNU General Public License v3.0',
-    author='Daniel Petro',
-    author_email='',
-    description='',
-    install_requires=[
-        'keras',
-        'sklearn',
-        'numpy',
-        'pandas',
-        'matplotlib',
-        'pillow', 'click'
-    ],
-    test_requires=[],
-    python_requires='>=3.6',    # We use some python 3.6+ only features, like ordered dicts by default and f-strings
-)