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/__init__.py

@@ -1,2 +1 @@
 from eyeballer.eyeballer import EyeballModel  # noqa: F401
-from eyeballer.visualization import OcclusionGenerator  # noqa: F401

eyeballer/eyeballer.py

@@ -3,7 +3,6 @@ import click
 import csv
 
 from model import EyeballModel, DATA_LABELS
-from visualization import HeatMap
 
 
 @click.group(invoke_without_command=True)
@@ -35,16 +34,10 @@ def train(ctx, graphs, batchsize, epochs):
 
 @cli.command()
 @click.argument('screenshot')
-@click.option('--heatmap', default=False, is_flag=True, help="Create a heatmap graphfor the prediction")
 @click.pass_context
-def predict(ctx, screenshot, heatmap):
+def predict(ctx, screenshot):
     model = EyeballModel(**ctx.obj['model_kwargs'])
     results = model.predict(screenshot)
-
-    if heatmap:
-        # Generate a heatmap
-        HeatMap(screenshot, model).generate()
-
     if not results:
         print("Error: Input file does not exist")
     if len(results) == 1:

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()
@@ -172,24 +196,6 @@ class EyeballModel:
             plt.legend(['Train', 'Validation'], loc='upper left')
             plt.savefig(self.graphs_directory + "loss.png")
 
-    def predict_on_array(self, image):
-        """Predict the labels for a single screenshot
-
-        Keyword arguments:
-        image -- The numpy array of the image to classify
-        """
-        img = np.expand_dims(image, axis=0)
-        img = preprocess_input(img)
-
-        prediction = self.model.predict(img, batch_size=1)
-        result = dict()
-        result["filename"] = "custom-image"
-        result["custom404"] = prediction[0][0]
-        result["login"] = prediction[0][1]
-        result["homepage"] = prediction[0][2]
-        result["oldlooking"] = prediction[0][3]
-        return result
-
     def predict(self, path):
         """Predict the labels for a single file
 
@@ -265,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

eyeballer/visualization.py

@@ -1,127 +0,0 @@
-import numpy as np
-import math
-from keras.preprocessing import image
-from copy import deepcopy
-from matplotlib import pyplot as plt
-from model import DATA_LABELS
-
-
-class HeatMap():
-    def __init__(self, screenshot_path, model, boxsize=28, step=7):
-        self.model = model
-
-        img = image.load_img(screenshot_path, target_size=(self.model.image_width, self.model.image_height))
-        img = image.img_to_array(img)
-
-        self.screenshot = img
-        self.boxsize = boxsize
-        self.step = step
-        self.x = 0
-        self.y = 0
-        self.screenshot_path = screenshot_path
-        self.gamma = 3
-
-    def generate(self, output_file="heatmap.png"):
-        """ Make a single heatmap image and return it """
-        heatmaps = []
-        labels = []
-        results = self.model.predict(self.screenshot_path)
-        results = results[0]
-        for label in DATA_LABELS:
-            boxsize = self.boxsize
-            # Ignore this label if it didn't predict positively (true label)
-            if results[label] > 0.5:
-                worst_score = 1
-                while worst_score > 0.5:
-                    heatmap, worst_score = self._get_heatmap(label, boxsize)
-                    if worst_score > 0.5:
-                        boxsize += 28
-                        print(f"Didn't get a good image for {label}. Trying again with a bigger boxsize: {boxsize}")
-                    else:
-                        heatmaps.append(heatmap)
-                        labels.append(label)
-
-        plt.figure()
-        screenshot_image = image.load_img(self.screenshot_path, target_size=(self.model.image_width, self.model.image_width))
-        _, subplots = plt.subplots(1, len(heatmaps))
-        # This is a little janky, but matplotlib returns a list above if there's multiple subplots, and just a single subplot if there's only one
-        #   it would have been easier if it was always a list. But alas
-        if not heatmaps:
-            print("No heatmap made. The image did not have a True classification for any label")
-            return
-        if len(heatmaps) > 1:
-            for i, heatmap in enumerate(heatmaps):
-                subplots[i].imshow(screenshot_image, cmap='binary', interpolation='none')
-                subplots[i].imshow(heatmap, alpha=0.5, interpolation='none')
-                subplots[i].set_title(labels[i])
-        else:
-            subplots.imshow(screenshot_image, cmap='binary', interpolation='none')
-            subplots.imshow(heatmap, alpha=0.5, interpolation='none')
-            subplots.set_title(labels[0])
-        plt.savefig(output_file)
-        print(f"Heatmap image written to: {output_file}")
-
-    def _get_heatmap(self, label, boxsize):
-        worst_score = 1
-        self.x, self.y = 0, 0
-
-        heatmap = np.ones((self.model.image_width, self.model.image_height))
-        heatmap *= 255
-        while True:
-            # Occlude an image
-            occluded_image, x, y = self._occlude(boxsize)
-            if occluded_image is not None:
-                # Score the new occluded image
-                results = self.model.predict_on_array(occluded_image)
-                score = results[label]  # TODO
-                worst_score = min(worst_score, score)
-                # Scale the score up to pixel values
-                score = math.floor(score * 255)
-                # Anneal the scores as you go outward into the occlusion block
-                new_scores = np.ones((boxsize, boxsize))
-                new_scores *= score
-                new_scores = self._gamma_anneal(new_scores)
-
-                # For each pixel, update only where the new score is lower
-                original_scores = heatmap[x: x+boxsize, y: y+boxsize]
-
-                target_area = np.where(new_scores < original_scores, new_scores, original_scores)
-                heatmap[x: x+boxsize, y: y+boxsize] = target_area
-            else:
-                break
-
-        return heatmap, worst_score
-
-    def _gamma_anneal(self, occlusion_area):
-        """ Get the annealed score based on centeredness of the occlusion area """
-        width, height = occlusion_area.shape
-
-        center_x = (width-1) / 2
-        center_y = (height-1) / 2
-
-        annealed_area = np.zeros(occlusion_area.shape)
-
-        for (x, y), score in np.ndenumerate(occlusion_area):
-            distance = math.sqrt((center_x - x)**2 + (center_y - y)**2)
-            annealed_area[x, y] = min(255, score + (self.gamma * distance))
-
-        return annealed_area
-
-    def _occlude(self, boxsize):
-        ''' Return a single occluded image and its location '''
-        if self.x + boxsize > self.screenshot.shape[0]:
-            return None, None, None
-
-        retImg = np.copy(self.screenshot)
-        retImg[self.x:self.x+boxsize, self.y:self.y+boxsize] = 0.0
-
-        old_i = deepcopy(self.x)
-        old_j = deepcopy(self.y)
-
-        # update indices
-        self.y = self.y + self.step
-        if self.y+boxsize > self.screenshot.shape[1]:  # reached end
-            self.y = 0  # reset j
-            self.x = self.x + self.step  # go to next row
-
-        return retImg, old_i, old_j

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
-)