Fixed progressbar.progressbar issue

- All prediction tests

.gitignore

@@ -113,3 +113,8 @@ venv.bak/
 *.model
 *.csv
 *.txt
+results.html
+!requirements.txt
+!requirements-gpu.txt
+!tests/models/*
+!tests/data/*

README.md

@@ -0,0 +1,47 @@
+# Eyeballer
+Give those screenshots of yours a quick eyeballing.
+
+Eyeballer is meant for large-scope network penetration tests where you need to find "interesting" targets from a huge set of web-based hosts. Go ahead and use your favorite screenshotting tool like normal (EyeWitness or GoWitness) and then run them through Eyeballer to tell you what's likely to contain vulnerabilities, and what isn't.
+
+## Setup
+
+Download required packages on pip:
+```
+sudo pip3 install -r requirements.txt
+```
+
+Or if you want GPU support:
+```
+sudo pip3 install -r requirements-gpu.txt
+```
+
+## Predicting Labels
+To eyeball some screenshots, just run the "predict" mode:
+
+```
+eyeballer.py --weights YOUR_WEIGHTS.h5 predict YOUR_FILE.png
+```
+
+Or for a whole directory of files:
+
+```
+eyeballer.py --weights YOUR_WEIGHTS.h5 predict PATH_TO/YOUR_FILES/
+```
+
+## Training
+To train a new model, run:
+```
+eyeballer.py train
+```
+
+You'll want a machine with a good GPU for this to run in a reasonable amount of time. Setting that up is outside the scope of this readme, however.
+
+This will output a new model file (weights.h5 by default).
+
+## Evaluation
+
+You just trained a new model, cool! Let's see how well it performs against some images it's never seen before, across a variety of metrics:
+
+```
+eyeballer.py --weights YOUR_WEIGHTS.h5 evaluate
+```

eyeballer.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+
+import click
+import csv
+
+from eyeballer.model import EyeballModel, DATA_LABELS
+from eyeballer.visualization import HeatMap
+from jinja2 import Template
+
+
+@click.group(invoke_without_command=True)
+@click.option('--weights', default=None, type=click.Path(), help="Weights file for input/output")
+@click.option('--summary/--no-summary', default=False, help="Print model summary at start")
+@click.option('--seed', default=None, type=int, help="RNG seed for data shuffling and transformations, defaults to random value")
+@click.pass_context
+def cli(ctx, weights, summary, seed):
+    model_kwargs = {"weights_file": weights,
+                    "print_summary": summary,
+                    "seed": seed}
+
+    #  pass the model to subcommands
+    ctx.ensure_object(dict)
+    # We only pass the kwargs so we can be lazy and make the model later after the subcommand cli is parsed. This
+    # way, the user doesn't have to wait for tensorflow if they are just calling --help on a subcommand.
+    ctx.obj['model_kwargs'] = model_kwargs
+
+
+@cli.command()
+@click.option('--graphs/--no-graphs', default=False, help="Save accuracy and loss graphs to file")
+@click.option('--epochs', default=20, type=int, help="Number of epochs")  # TODO better help string
+@click.option('--batchsize', default=32, type=int, help="Batch size")  # TODO better help string
+@click.pass_context
+def train(ctx, graphs, batchsize, epochs):
+    model = EyeballModel(**ctx.obj['model_kwargs'])
+    model.train(print_graphs=graphs, batch_size=batchsize, epochs=epochs)
+
+
+@cli.command()
+@click.argument('screenshot')
+@click.option('--heatmap', default=False, is_flag=True, help="Create a heatmap graphfor the prediction")
+@click.option('--threshold', default=.5, type=float, help="Threshold confidence for labeling")
+@click.pass_context
+def predict(ctx, screenshot, heatmap, threshold):
+    model = EyeballModel(**ctx.obj['model_kwargs'])
+    results = model.predict(screenshot)
+
+    if heatmap:
+        # Generate a heatmap
+        HeatMap(screenshot, model, threshold).generate()
+
+    if not results:
+        print("Error: Input file does not exist")
+    if len(results) == 1:
+        print(results)
+    else:
+        with open("results.csv", "w", newline="") as csvfile:
+            fieldnames = ["filename", "custom404", "login", "homepage", "oldlooking"]
+            labelwriter = csv.DictWriter(csvfile, fieldnames=fieldnames)
+            labelwriter.writeheader()
+            labelwriter.writerows(results)
+
+        print("Output written to results.csv")
+        buildHTML(processResults(results, threshold))
+        print("HTML written to results.html")
+
+
+def processResults(results, threshold):
+    '''Filter the initial results dictionary and reformat it for use in JS.
+
+        Keyword arguments:
+        results -- dictionary output from predict function
+
+    '''
+    jsResults = {}
+
+    for result in results:
+        positiveTags = []
+
+        for label, label_info in result.items():
+            if (label == 'filename'):
+                pass
+            elif label_info > threshold:
+                positiveTags.append(label)
+
+        jsResults[result['filename']] = positiveTags
+    return(jsResults)
+
+
+def buildHTML(jsResults):
+    '''Build HTML around the JS Dictionary that is passed from processResults.
+
+        Keyword arguments:
+        jsResults -- dictionary output from processResults function
+    '''
+    html_output = ""
+    with open("prediction_output_template.html") as template_file:
+        template = Template(template_file.read())
+        html_output = template.render(jsResults=jsResults)
+
+    with open('results.html', 'w') as file:
+        file.write(html_output)
+
+
+def pretty_print_evaluation(results):
+    """Print a human-readable summary of the evaluation"""
+    # We use 4.2% to handle all the way from "  0.00%" (7chars) to "100.00%" (7chars)
+    for label in DATA_LABELS:
+        print("{} Precision Score: {:4.2%}".format(label, results[label]['precision']))
+        print("{} Recall Score: {:4.2%}".format(label, results[label]['recall']))
+    print("'None of the above' Precision: {:4.2%}".format(results['none_of_the_above_precision']))
+    print("'None of the above' Recall: {:4.2%}".format(results['none_of_the_above_recall']))
+    print("All or nothing Accuracy: {:4.2%}".format(results['all_or_nothing_accuracy']))
+    print("Overall Binary Accuracy: {:4.2%}".format(results['total_binary_accuracy']))
+    print("Top 10 worst predictions: {}".format(results['top_10_worst'][1]))
+
+
+@cli.command()
+@click.option('--threshold', default=.5, type=float, help="Threshold confidence for labeling")
+@click.pass_context
+def evaluate(ctx, threshold):
+    model = EyeballModel(**ctx.obj['model_kwargs'])
+    results = model.evaluate(threshold)
+    pretty_print_evaluation(results)
+
+
+if __name__ == '__main__':
+    cli()

eyeballer/__init__.py

@@ -1 +0,0 @@
-from eyeballer.eyeballer import EyeballModel  # noqa: F401

eyeballer/augmentation.py

@@ -0,0 +1,18 @@
+import numpy as np
+
+from Augmentor.Operations import Operation
+from keras.applications.mobilenet import preprocess_input
+
+
+class EyeballerAugmentation(Operation):
+    def __init__(self, probability=1):
+        Operation.__init__(self, probability)
+
+    # Class must implement the perform_operation method
+    def perform_operation(self, images):
+        return_list = []
+        for image in images:
+            image_array = np.array(image).astype('uint8')
+            image_array = preprocess_input(image_array)
+            return_list.append(image_array)
+        return return_list

eyeballer/eyeballer.py

@@ -1,62 +0,0 @@
-#!/usr/bin/env python3
-import click
-from model import EyeballModel, DATA_LABELS
-
-
-@click.group(invoke_without_command=True)
-@click.option('--weights', default=None, type=click.Path(), help="Weights file for input/output")
-@click.option('--summary/--no-summary', default=False, help="Print model summary at start")
-@click.option('--seed', default=None, type=int, help="RNG seed for data shuffling and transformations, defaults to random value")
-@click.pass_context
-def cli(ctx, weights, summary, seed):
-    model_kwargs = {"weights_file": weights,
-                    "print_summary": summary,
-                    "seed": seed}
-
-    #  pass the model to subcommands
-    ctx.ensure_object(dict)
-    # We only pass the kwargs so we can be lazy and make the model later after the subcommand cli is parsed. This
-    # way, the user doesn't have to wait for tensorflow if they are just calling --help on a subcommand.
-    ctx.obj['model_kwargs'] = model_kwargs
-
-
-@cli.command()
-@click.option('--graphs/--no-graphs', default=False, help="Save accuracy and loss graphs to file")
-@click.option('--epochs', default=20, type=int, help="Number of epochs")  # TODO better help string
-@click.option('--batchsize', default=32, type=int, help="Batch size")  # TODO better help string
-@click.pass_context
-def train(ctx, graphs, batchsize, epochs):
-    model = EyeballModel(**ctx.obj['model_kwargs'])
-    model.train(print_graphs=graphs, batch_size=batchsize, epochs=epochs)
-
-
-@cli.command()
-@click.argument('screenshot', nargs=-1)  # Nargs=-1 means you can pass as many in as you want
-@click.pass_context
-def predict(ctx, screenshot):
-    model = EyeballModel(**ctx.obj['model_kwargs'])
-    for fn in screenshot:
-        model.predict(fn)
-
-
-def pretty_print_evaluation(results):
-    """Print a human-readable summary of the evaluation"""
-    # We use 4.2% to handle all the way from "  0.00%" (7chars) to "100.00%" (7chars)
-    for label in DATA_LABELS:
-        print(f"{label} Precision Score: {results[label]['precision']:4.2%}")
-        print(f"{label} Precision Score: {results[label]['recall']:4.2%}")
-    print(f"Overall Binary Accuracy: {results['total_binary_accuracy']:4.2%}")
-    print(f"All or nothing accuracy: {results['all_or_nothing_accuracy']:4.2%}")
-
-
-@cli.command()
-@click.option('--threshold', default=.5, type=float, help="Threshold confidence for labeling")
-@click.pass_context
-def evaluate(ctx, threshold):
-    model = EyeballModel(**ctx.obj['model_kwargs'])
-    results = model.evaluate(threshold)
-    pretty_print_evaluation(results)
-
-
-if __name__ == '__main__':
-    cli()

eyeballer/model.py

@@ -1,24 +1,23 @@
 import os
 import random
-import time
+import sys
+import progressbar
 
+# 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
+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 import image
-from keras.preprocessing.image import ImageDataGenerator
-from sklearn.metrics import classification_report, accuracy_score, jaccard_similarity_score
+from tensorflow.keras.applications.mobilenet import preprocess_input
+from sklearn.metrics import classification_report, accuracy_score, hamming_loss
+from eyeballer.augmentation import EyeballerAugmentation
 
-DATA_LABELS = ["custom404", "login", "homepage"]
+DATA_LABELS = ["custom404", "login", "homepage", "oldlooking"]
 
 
 class EyeballModel:
@@ -26,12 +25,13 @@ class EyeballModel:
 
     Contains high-level functions for training, evaluating, and predicting.
     """
-    checkpoint_file = "weights-snapshot.h5"
-    image_dir = "images_224x224/"
+    graphs_directory = "graphs/"
+    checkpoint_file = "weights.h5"
+    image_dir = "images/"
     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:
@@ -39,22 +39,23 @@ 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(3, 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.MobileNet(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:
-        #    layer.trainable = False
+        self.model.compile(optimizer=tf.keras.optimizers.Adam(0.0005),
+                           loss="binary_crossentropy",
+                           metrics=["accuracy"])
 
-        adam = Adam(lr=0.0005)
-        self.model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
+        # for layer in pretrained_layer.layers:
+        #     layer.trainable = False
 
         if print_summary:
             print(self.model.summary())
@@ -71,15 +72,35 @@ class EyeballModel:
             self.random_seed = True
             self.seed = random.randint(0, 999999)
             print("No seed set, ", end='')
-        print(f"using seed: {self.seed}")
+        print("using seed: {}".format(self.seed))
         random.seed(self.seed)
         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:
+                print("ERROR: Unable to open weights file '{}'".foramt(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()
+        augmentor.set_seed(self.seed)
+        augmentor.zoom(probability=0.75, min_factor=0.8, max_factor=1.2)
+        augmentor.random_color(probability=0.75, min_factor=0.5, max_factor=1.0)
+        augmentor.random_contrast(probability=0.75, min_factor=0.8, max_factor=1.0)
+        augmentor.random_brightness(probability=0.75, min_factor=0.8, max_factor=1.2)
+        augmentor.random_erasing(probability=0.75, rectangle_area=0.15)
+
+        # Finalizes the augmentation with a custom operation to prepare the image for the specific pretrained model we're using
+        training_augmentation = EyeballerAugmentation()
+        augmentor.add_operation(training_augmentation)
+        self.preprocess_training_function = augmentor.keras_preprocess_func()
 
     def train(self, epochs=20, batch_size=32, print_graphs=False):
         """Train the model, making a new weights file at each successfull checkpoint. You'll probably need a GPU for this to realistically run.
@@ -91,15 +112,10 @@ class EyeballModel:
         """
         print("Training with seed: " + str(self.seed))
 
-        # Data augmentation
-        data_generator = ImageDataGenerator(
-            preprocessing_function=preprocess_input,
-            shear_range=0.0,
-            zoom_range=0.2,
-            brightness_range=(0.7, 1.3),
-            samplewise_center=True,
+        data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
+            preprocessing_function=self.preprocess_training_function,
             validation_split=0.2,
-            horizontal_flip=False)
+            samplewise_center=True)
 
         training_generator = data_generator.flow_from_dataframe(
             self.training_labels,
@@ -125,8 +141,12 @@ 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,
@@ -138,6 +158,8 @@ class EyeballModel:
             verbose=1)
 
         if print_graphs:
+            if not os.path.exists(self.graphs_directory):
+                os.makedirs(self.graphs_directory)
             # Plot training & validation accuracy values
             plt.plot(history.history['acc'])
             plt.plot(history.history['val_acc'])
@@ -145,7 +167,7 @@ class EyeballModel:
             plt.ylabel('Accuracy')
             plt.xlabel('Epoch')
             plt.legend(['Train', 'Validation'], loc='upper left')
-            plt.savefig("accuracy.png")
+            plt.savefig(self.graphs_directory + "accuracy.png")
             plt.clf()
             plt.cla()
             plt.close()
@@ -157,28 +179,68 @@ class EyeballModel:
             plt.ylabel('Loss')
             plt.xlabel('Epoch')
             plt.legend(['Train', 'Validation'], loc='upper left')
-            plt.savefig("loss.png")
+            plt.savefig(self.graphs_directory + "loss.png")
 
-    def predict(self, filename):
-        """Predict the labels for a single file
+    def predict_on_array(self, image):
+        """Predict the labels for a single screenshot
 
         Keyword arguments:
-        filename -- The path to the file that we'll be evaluating
+        image -- The numpy array of the image to classify
         """
-        # Load the image into memory
-        img = image.load_img(filename, target_size=(self.image_width, self.image_height))
-        img = image.img_to_array(img)
-        img = np.expand_dims(img, axis=0)
+        img = np.expand_dims(image, axis=0)
         img = preprocess_input(img)
 
-        time_start = time.time()
         prediction = self.model.predict(img, batch_size=1)
-        time_end = time.time()
-        print("Predictions:")
-        print("\tCustom 404: " + str(round(prediction[0][0] * 100, 2)) + "%")
-        print("\tLogin Page: " + str(round(prediction[0][1] * 100, 2)) + "%")
-        print("\tHome Page: " + str(round(prediction[0][2] * 100, 2)) + "%")
-        print("Prediction Took (seconds): ", time_end - time_start)
+        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, threshold=0.5):
+        """Predict the labels for a single file or directory of files
+
+        Keyword arguments:
+        path -- The path to the file(s) that we'll be evaluating.
+        """
+        # Is this a single file, or a directory?
+        screenshots = []
+        if os.path.isfile(path):
+            screenshots = [path]
+        elif os.path.isdir(path):
+            screenshots = os.listdir(path)
+            screenshots = [os.path.join(path, s) for s in screenshots]
+        else:
+            raise FileNotFoundError
+
+        results = []
+        pbar = progressbar.ProgressBar()
+        for screenshot in pbar(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
+            except OSError:
+                print("\nWARN: Skipping empty or corrupt file: ", screenshot)
+                continue
+
+            prediction = self.model.predict(img, batch_size=1)
+            result = dict()
+            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
 
     def evaluate(self, threshold=0.5):
         """Evaluate performance against the persistent evaluation data set
@@ -187,7 +249,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,
@@ -202,14 +264,10 @@ class EyeballModel:
         if not self.random_seed:
             print("Using validation set...")
             # Data augmentation
-            data_generator = ImageDataGenerator(
-                preprocessing_function=preprocess_input,
-                shear_range=0.0,
-                zoom_range=0.2,
-                brightness_range=(0.7, 1.3),
+            data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
+                preprocessing_function=self.preprocess_training_function,
                 samplewise_center=True,
-                validation_split=0.2,
-                horizontal_flip=False)
+                validation_split=0.2)
             evaluation_generator = data_generator.flow_from_dataframe(
                 self.training_labels,
                 directory=self.image_dir,
@@ -224,13 +282,104 @@ class EyeballModel:
         else:
             print("Using evaluation set...")
 
-        predictions = self.model.predict_generator(evaluation_generator, verbose=1, steps=len(evaluation_generator))
+        predictions = self.model.predict_generator(
+            evaluation_generator,
+            verbose=1,
+            steps=len(evaluation_generator))
+
+        self._save_prediction_histograms(predictions)
+        print("predictions: ", type(predictions))
         predictions = predictions > threshold
-        ground_truth = evaluation_generator.data
+        ground_truth = self.evaluation_labels[DATA_LABELS].to_numpy()
+        filenames = self.evaluation_labels[["filename"]].to_numpy()
         stats = classification_report(ground_truth, predictions, target_names=DATA_LABELS, output_dict=True)
-        stats["total_binary_accuracy"] = jaccard_similarity_score(ground_truth, predictions)
+        stats["total_binary_accuracy"] = 1 - hamming_loss(ground_truth, predictions)
         stats["all_or_nothing_accuracy"] = accuracy_score(ground_truth, predictions)
+        stats["top_10_best"] = self._top_images(filenames, ground_truth, predictions, best=True)
+        stats["top_10_worst"] = self._top_images(filenames, ground_truth, predictions, best=False)
+        stats["none_of_the_above_recall"] = self._none_of_the_above_recall(ground_truth, predictions)
+        stats["none_of_the_above_precision"] = self._none_of_the_above_precision(ground_truth, predictions)
         return stats
 
-    def get_data_column(self, data_slice, data):
-        return np.reshape(np.array(data)[data_slice], len(data)).tolist()
+    def _none_of_the_above_recall(self, labels, predictions):
+        """Returns the recall score for the 'none of the above' images.
+        That means, all the images that don't have a category.
+        """
+        total_count = 0
+        correct_count = 0
+        for item in zip(labels.astype(int), predictions.astype(int)):
+            # Is this a none of the above?
+            if not item[0].any():
+                total_count += 1
+                if not item[1].any():
+                    correct_count += 1
+        if total_count == 0:
+            print("WARNING: None of the Above Recall is NaN")
+            return 0
+        return correct_count / total_count
+
+    def _none_of_the_above_precision(self, labels, predictions):
+        """Returns the precision score for the 'none of the above' images.
+        That means, all the images that don't have a category.
+        """
+        total_count = 0
+        correct_count = 0
+        for item in zip(labels.astype(int), predictions.astype(int)):
+            # Is this a none of the above prediction?
+            if not item[1].any():
+                total_count += 1
+                if not item[0].any():
+                    correct_count += 1
+        if total_count == 0:
+            print("WARNING: None of the Above Precision is NaN")
+            return 0
+        return correct_count / total_count
+
+    def _top_images(self, filenames, ground_truth, predictions, top_k=10, best=False):
+        """Collect top-k best or top-k worst predicted images
+
+        Keyword arguments:
+        ground_truth -- The correct labels
+        predictions -- The numpy array of predictions
+        top_k -- Top k elements
+        best -- True/False. Calculate either the best or worst images
+
+        :Return -- Tuple of top-k indicies and top-k filenames
+        """
+        true_labels = np.array(ground_truth).astype(float)
+        predictions = predictions.astype(float)
+
+        differences = np.abs(predictions - true_labels).sum(axis=1)
+        indicies = np.argsort(differences, axis=0)
+
+        top_file_list = []
+
+        if not best:
+            # Reverse numpy array
+            indicies = np.flipud(indicies)
+
+        for i in indicies[:top_k]:
+            top_file_list.append(filenames[i][0])
+
+        return indicies[:top_k], top_file_list
+
+    def _save_prediction_histograms(self, predictions, buckets=50):
+        """Saves a series of histogram screenshots
+
+            Keyword arguments:
+            predictions -- The numpy array of predicted labels, ranging from 0->1.
+            buckets -- The number of buckets to divide the data into. Default: 50
+            :Returns -- Nothing"
+        """
+        figure, axes = plt.subplots(nrows=len(DATA_LABELS))
+        for i, label in enumerate(DATA_LABELS):
+            axes[i].hist(predictions[:, i], buckets, alpha=.75)
+            axes[i].set_xlabel("Prediction")
+            axes[i].set_ylabel("Counts of Predictions")
+            axes[i].set_title(label)
+            axes[i].grid(True)
+        figure.set_size_inches(5, 3*len(DATA_LABELS))
+        figure.tight_layout()
+        if not os.path.exists(self.graphs_directory):
+            os.makedirs(self.graphs_directory)
+        plt.savefig(self.graphs_directory + "label_histograms.png")

eyeballer/visualization.py

@@ -0,0 +1,128 @@
+import numpy as np
+import math
+from keras.preprocessing import image
+from copy import deepcopy
+from matplotlib import pyplot as plt
+from eyeballer.model import DATA_LABELS
+
+
+class HeatMap():
+    def __init__(self, screenshot_path, model, threshold=0.5, 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
+        self.threshold = threshold
+
+    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] > self.threshold:
+                worst_score = 1
+                while worst_score > self.threshold:
+                    heatmap, worst_score = self._get_heatmap(label, boxsize)
+                    if worst_score > self.threshold:
+                        boxsize += 28
+                        print("Didn't get a good image for {}. Trying again with a bigger boxsize: {}".format(label, 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("Heatmap image written to: {}".format(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

prediction_output_template.html

@@ -0,0 +1,189 @@
+<!DOCTYPE HTML>
+<meta charset="utf-8"/>
+<html>
+<title>Eyeballer Results</title>
+<body>
+    <div id="btnsCont">
+      <button class="btn all" onclick="filterSelection('all')"> Show all</button>
+      <button class="btn" onclick="filterSelection('custom404')"> 404</button>
+      <button class="btn" onclick="filterSelection('login')"> Login</button>
+      <button class="btn" onclick="filterSelection('homepage')"> Home Page</button>
+      <button class="btn" onclick="filterSelection('oldlooking')"> Old Looking</button>
+        <select onchange="numOnPage=this.value; filterSelection([]);" style="float:right; margin:5px; margin-right:30px; height:30px; width:auto;">
+          <option value="25">25</option>
+          <option value="50">50</option>
+          <option value="100">100</option>
+          <option value="250">250</option>
+        </select>
+    </div>
+    <div id="btnsPage">
+    </div>
+    <div id="container">
+    </div>
+    <button onclick="topFunction()" id="topBtn" title="Go to top">Top</button>
+<script>
+    var results = {{jsResults}};
+    var selectedTags = new Array();
+    var btnContainer = document.getElementById("btnsCont");
+    var categoryBtns = btnContainer.getElementsByClassName("btn");
+    var numOnPage = 100;
+    var indexesToDisplay = [];
+    filterSelection([]);
+    for (var i = 0; i < categoryBtns.length; i++) {
+        categoryBtns[i].addEventListener("click", function () {
+            if (this.className.indexOf("active") > -1) {
+                this.className = this.className.replace(" active", "");
+            }
+            else {
+            this.className += " active";
+            }
+        });
+    }
+    function filterSelection(c) {
+        indexesToDisplay = [];
+        var currentIndex = -1;
+        if (c == "all") location.reload();
+        if (selectedTags.includes(c)!=true) selectedTags=selectedTags.concat(c);
+        else selectedTags.splice(selectedTags.indexOf(c),1);
+        for (var filename in results){
+                currentIndex++;
+                var tags = results[filename];
+                var isMissingTag = false;
+                for (var i in selectedTags){
+                    if (!tags.includes(selectedTags[i])){
+                        isMissingTag = true;
+                        break;
+                    }
+                }
+                if (isMissingTag === false){
+                    indexesToDisplay = indexesToDisplay.concat(currentIndex);
+                }
+            }
+        buildPageButtons(indexesToDisplay);
+        buildHTML(0, numOnPage);
+    }
+    function buildPageButtons(indexes){
+        numOfPages = Math.ceil(indexes.length/numOnPage);
+        var html = '';
+        for (var i = 0; i < numOfPages; i++){
+            html += ('<button class="btnpage" onclick="buildHTML(' + (i*numOnPage) + ', ' + ((i+1)*numOnPage) + '); activatePage(this);">' + (i+1) +'</button>');
+        }
+        document.getElementById('btnsPage').innerHTML = html;
+    }
+    function activatePage(pressedBtn){
+        var activeNow = document.getElementsByClassName('activePage')
+        for (var i = 0; i < activeNow.length; i++){
+            activeNow[i].className = activeNow[i].className.replace(" activePage", "");
+        }
+        pressedBtn.className += ' activePage';
+    }
+    function buildHTML(rangeMin, rangeMax){
+        var html = '';
+        if (rangeMax > indexesToDisplay.length) rangeMax=indexesToDisplay.length;
+        for (var i = rangeMin; i < rangeMax; i++){
+            index=indexesToDisplay[i];
+            filename = Object.keys(results)[index];
+            html += ('<div class="filterDiv"><a href=' + filename + ' target="_blank"> ' + results[filename].toString().replace(/,/g, ' ') + '</a><img class="res_pic" src="' + filename + '"></div>');
+        }
+        document.getElementById('container').innerHTML = html;
+    }
+    window.onscroll = function() {scrollFunction()};
+    function scrollFunction() {
+      if (document.body.scrollTop > 20 || document.documentElement.scrollTop > 20) {
+        document.getElementById("topBtn").style.display = "block";
+      } else {
+        document.getElementById("topBtn").style.display = "none";
+      }
+    }
+    function topFunction() {
+      document.body.scrollTop = 0; // For Safari
+      document.documentElement.scrollTop = 0; // For Chrome, Firefox, IE and Opera
+    }
+</script>
+    <style>
+    #topBtn {
+      display: none; /* Hidden by default */
+      position: fixed; /* Fixed/sticky position */
+      bottom: 20px; /* Place the button at the bottom of the page */
+      right: 30px; /* Place the button 30px from the right */
+      z-index: 99; /* Make sure it does not overlap */
+      border: none; /* Remove borders */
+      outline: none; /* Remove outline */
+      background-color: #193b72; /* Set a background color */
+      color: white; /* Text color */
+      cursor: pointer; /* Add a mouse pointer on hover */
+      padding: 15px; /* Some padding */
+      border-radius: 10px; /* Rounded corners */
+      font-size: 18px; /* Increase font size */
+      font-weight:bold;
+      border:solid;
+    }
+    #topBtn:hover {
+      background-color: #555;
+    }
+    .container {
+    overflow: hidden;
+    }
+    #btnsPage {
+    margin:10px;
+    }
+    .btnpage{
+        width:40px;
+        height:30px;
+        background-color: #f1f1f1;
+        outline: none;
+    }
+    .res_pic{
+    width:100%;
+    height:auto;
+    }
+    a{
+    font-weight:600;
+    padding:3px;
+    text-decoration: none;
+    color:#ffffff;
+    font-size:16px;
+    font-family:Sans-serif;
+    }
+    .name{
+    float:left;
+    }
+    .filterDiv {
+      float: left;
+      background-color: #777;
+      color: #ffffff;
+      width: 22%;
+      border-style:solid;
+      border-color:#666;
+      text-align: center;
+      margin: 10px;
+      line-height: 30px;
+      display: block;
+    }
+    .btn {
+      border: none;
+      outline: none;
+      padding: 12px 16px;
+      margin-left:10px;
+      background-color: #f1f1f1;
+      cursor: pointer;
+    }
+    .btn:hover {
+      background-color: #ddd;
+    }
+    .active {
+      background-color: #666;
+      color: white;
+    }
+    .activePage{
+        border: solid;
+        border-color:#666666;
+        color: black;
+    }
+    .btn.all {
+      background-color: #193b72;
+      color: white;
+    }
+    </style>
+</body>
+</html>

requirements-gpu.txt

@@ -0,0 +1,10 @@
+Augmentor
+click
+matplotlib
+numpy
+pandas
+pillow
+sklearn
+tensorflow-gpu==2.0.0-beta0
+jinja2
+progressbar2

requirements.txt

@@ -0,0 +1,10 @@
+Augmentor
+click
+matplotlib
+numpy
+pandas
+pillow
+sklearn
+tensorflow==2.0.0-beta0
+jinja2
+progressbar2

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

tests/EvaluateTest.py

@@ -0,0 +1 @@
+#TODO

tests/PredictTest.py

@@ -0,0 +1,128 @@
+import errno
+import os
+import unittest
+import sys
+
+from eyeballer.model import EyeballModel
+
+
+class PredictTest(unittest.TestCase):
+    def setUp(self):
+
+        class DummyFile(object):
+            def write(self, x): pass
+            def flush(self): pass
+
+        sys.stdout = DummyFile()
+        weights_file = "tests/models/test_weights.h5"
+        if not os.path.isfile(weights_file):
+            print("Error: Symlink the latest weights file to " + weights_file)
+            raise FileNotFoundError(
+                errno.ENOENT,
+                os.strerror(errno.ENOENT),
+                weights_file)
+
+        model_kwargs = {
+            "weights_file": weights_file,
+            "print_summary": False,
+            "seed": None
+        }
+        self.model = EyeballModel(**model_kwargs)
+
+    def test_different_seed_predict(self):
+        model_kwargs = {
+            "weights_file": None,
+            "print_summary": False,
+            "seed": 12345678
+        }
+        same_seed_model = EyeballModel(**model_kwargs)
+
+        screenshot = "tests/data/404.png"
+
+        results_one = self.model.predict(screenshot)
+        results_two = same_seed_model.predict(screenshot)
+
+        self.assertNotEqual(results_one, results_two)
+
+    def test_same_seed_predict(self):
+        model_kwargs = {
+            "weights_file": None,
+            "print_summary": False,
+            "seed": 12345678
+        }
+        same_seed_model = EyeballModel(**model_kwargs)
+
+        screenshot = "tests/data/404.png"
+
+        results_one = same_seed_model.predict(screenshot)
+        results_two = same_seed_model.predict(screenshot)
+
+        self.assertEqual(results_one, results_two)
+
+    def test_predict_custom404(self):
+        screenshot = "tests/data/404.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["custom404"], 0.5)
+
+    def test_predict_not_custom404(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertLess(results["custom404"], 0.5)
+
+    def test_predict_login(self):
+        screenshot = "tests/data/login.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["login"], 0.5)
+
+    def test_predict_not_login(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        print(screenshot, results)
+        self.assertLess(results["login"], 0.5)
+
+    def test_predict_homepage(self):
+        screenshot = "tests/data/homepage.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["homepage"], 0.5)
+
+    def test_predict_not_homepage(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertLess(results["homepage"], 0.5)
+
+    def test_predict_oldlooking(self):
+        screenshot = "tests/data/old-looking.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["oldlooking"], 0.5)
+
+    def test_predict_not_oldlooking(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertLess(results["oldlooking"], 0.5)
+
+    def test_file_doesnt_exist(self):
+        screenshot = "tests/data/doesnotexist.png"
+        try:
+            self.model.predict(screenshot)[0]
+            self.fail("FileNotFoundError was expected but not found")
+        except FileNotFoundError:
+            pass
+
+    def test_folder(self):
+        screenshots = "tests/data/"
+        results = self.model.predict(screenshots)
+        self.assertEqual(len(results), 7)
+
+    def test_file_is_empty(self):
+        """
+        We're just testing that it doesn't crash, basically
+        """
+        screenshot = "tests/data/empty.png"
+        self.model.predict(screenshot)
+
+    def test_file_is_invalid(self):
+        """
+        We're just testing that it doesn't crash, basically
+        """
+        screenshot = "tests/data/invalid.png"
+        self.model.predict(screenshot)

tests/TrainTest.py

@@ -0,0 +1 @@
+#TODO

tests/__init__.py

@@ -0,0 +1 @@
+from tests.PredictTest import *

tests/data/invalid.png

@@ -0,0 +1 @@
+THIS TEXT IS NOT A VALID PNG

utils/labelbox_to_labels.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+
+import csv
+import random
+import json
+
+with open("newlabels.csv", "w", newline="") as csvfile:
+    # Open the old labels file
+    with open("labelbox.csv", newline="") as oldfile:
+        # Get the header labels
+        csvreader = csv.DictReader(oldfile)
+        fieldnames = next(csvreader)
+        labelwriter = csv.DictWriter(csvfile, fieldnames=["filename", "login", "custom404", "homepage", "oldlooking", "evaluation"])
+        labelwriter.writeheader()
+
+        # Loop through the file
+        rows = []
+        for row in csvreader:
+            filename = row["External ID"]
+            print(row["Label"])
+            labelstring = row["Label"]
+            if row["Label"] == "Skip":
+                labelstring = '{"imageclassification":[]}'
+            labels = json.loads(labelstring)
+
+            newrow = dict()
+            newrow["filename"] = filename
+            newrow["oldlooking"] = True
+            newrow["login"] = "loginpage" in labels["imageclassification"]
+            newrow["homepage"] = "homepage" in labels["imageclassification"]
+            newrow["custom404"] = "custom404"in labels["imageclassification"]
+
+            newrow["evaluation"] = random.random() > 0.8
+
+            rows.append(newrow)
+        labelwriter.writerows(rows)
+print("Made new labels file: newlabels.csv")