Evaluator can be stopped early and output results

Callback object contains intermediate results from evaluator function and
will output results if user sends a keyboard interrupt.

.gitignore

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

eyeballer/model.py

@@ -1,7 +1,11 @@
 import os
 import random
 
+# Prevent Tkinter Dependency
+import matplotlib
+matplotlib.use('agg')
 import matplotlib.pyplot as plt
+
 import numpy as np
 import pandas as pd
 import Augmentor
@@ -18,10 +22,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.
 
@@ -247,7 +264,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:
+            output = 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
-)