Added a GPU check for NVIDIA in Makefile

This fix will only work on Linux. We would need to use 'ioreg' for Mac

.gitignore

@@ -113,10 +113,4 @@ venv.bak/
 *.model
 *.csv
 *.txt
-results.html
 !requirements.txt
-!requirements-gpu.txt
-!tests/models/*
-!tests/data/*
-!docs/*
-logs/*

Makefile

@@ -0,0 +1,99 @@
+###############################################################################
+# Parts of this are borrowed from the HTTPie Makefile. Thanks!
+###############################################################################
+
+REQUIREMENTS="requirements-dev.txt"
+TAG="\n\n\033[0;32m\#\#\# "
+END=" \#\#\# \033[0m\n"
+
+PIP=$(shell which pip3)
+CONDA=$(shell which conda)
+CONDA_HOME="$(HOME)/anaconda3"
+VENV="eyeballer_venv"
+
+SHELL := /bin/bash
+.ONE_SHELL:
+
+all: test
+
+# Target-specific assignments
+conda: PIP := "$(CONDA_HOME)/envs/$(VENV)/bin/pip"
+conda init: uninstall-eyeballer
+
+ifeq ($(strip $(MAKECMDGOALS)),conda)
+	@echo -e $(TAG)Conda detected at: $(CONDA)$(END)
+	conda create -yn $(VENV) | true
+
+	source $(CONDA_HOME)/etc/profile.d/conda.sh; \
+	conda activate $(VENV); \
+	conda install -y pip
+endif
+	@echo -e $(TAG)Installing dev requirements$(END)
+	$(PIP) install --upgrade -r $(REQUIREMENTS)
+
+	@echo -e $(TAG)Installing Eyeballer$(END)
+ifeq ($(shell lspci | grep NVIDIA),NVIDIA))
+	$(PIP) install --upgrade --editable .[tf_gpu]
+else
+	$(PIP) install --upgrade --editable .[tf]
+endif
+
+ifeq ($(strip $(MAKECMDGOALS)),conda)
+	@echo -e $(TAG)Activate virtualenv to run eyeballer: $(END)
+	@echo conda activate $(VENV)
+endif
+
+	@echo
+
+clean:
+	@echo -e $(TAG)Cleaning up$(END)
+	rm -rf *.egg dist build .cache .pytest_cache eyeballer.egg-info
+	find . -name '__pycache__' -delete -print -o -name '*.pyc' -delete -print
+ifneq ($(strip $(CONDA)),)
+	conda env remove -y --name eyeballer_venv | true
+endif
+	@echo
+
+
+###############################################################################
+# Testing
+###############################################################################
+
+
+test:
+	@echo -e $(TAG)Running tests on the current Python interpreter with coverage $(END)
+	python -m unittest tests
+	@echo
+
+
+# test-all is meant to test everything — even this Makefile
+test-all: uninstall-all clean init test pycodestyle
+	@echo
+
+
+pycodestyle:
+	which pycodestyle || $(PIP) install pycodestyle
+	pycodestyle
+	@echo
+
+
+###############################################################################
+# Uninstalling
+###############################################################################
+
+uninstall-eyeballer:
+	@echo using pip3 located at: $(PIP)
+	@echo -e $(TAG)Uninstalling eyeballer$(END)
+	- $(PIP) uninstall --yes eyeballer &2>/dev/null
+
+	@echo "Verifying…"
+	cd .. && ! python -m eyeballer --version &2>/dev/null
+
+	@echo "Done"
+	@echo
+
+
+uninstall-all: uninstall-eyeballer
+
+	@echo -e $(TAG)Uninstalling development requirements$(END)
+	- $(PIP) uninstall --yes -r $(REQUIREMENTS)

README.md

@@ -3,77 +3,44 @@ 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.
 
+## Installation
 
-### Example Labels
-
-| Old-Looking Sites | Login Pages |
-| ------ |:-----:|
-| ![Sample Old-looking Page](/docs/old-looking.png) | ![Sample Login Page](/docs/login.png) |
-
-| Webapp | Custom 404's |
-| ------ |:-----:|
-| ![Sample Webapp](/docs/homepage.png) | ![Sample Custom 404](/docs/404.png) |
-
-## What the Labels Mean
-
-**Old-Looking Sites**
-Blocky frames, broken CSS, that certain "je ne sais quoi" of a website that looks like it was designed in the early 2000's. You know it when you see it. Old websites aren't just ugly, they're also typically super vulnerable. When you're looking to hack into something, these websites are a gold mine.
-
-**Login Pages**
-Login pages are valuable to pen testing, they indicate that there's additional functionality you don't currently have access to. It also means there's a simple follow-up process of credential enumeration attacks. You might think that you can set a simple heuristic to find login pages, but in practice it's really hard. Modern sites don't just use a simple input tag we can grep for.
-
-**Webapp**
-This tells you that there is a larger group of pages and functionality available here that can serve as surface area to attack. This is in contrast to a simple login page, with no other functionality. Or a default IIS landing page which has no other functionality. This label should indicate to you that there is a web application here to attack.
-
-**Custom 404**
-Modern sites love to have cutesy custom 404 pages with pictures of broken robots or sad looking dogs. Unfortunately, they also love to return HTTP 200 response codes while they do it. More often, the "404" page doesn't even contain the text "404" in it. These pages are typically uninteresting, despite having a lot going on visually, and Eyeballer can help you sift them out.
-
-## Setup
-
-Download required packages on pip:
+To install in a python venv (Recommended):
 ```
-sudo pip3 install -r requirements.txt
+python3 -m venv venv
+source venv/bin/activate
+make init
 ```
 
-Or if you want GPU support:
+To install for Conda (automatically creates conda venv):
+
 ```
-sudo pip3 install -r requirements-gpu.txt
+make conda
 ```
 
-**NOTE**: Setting up a GPU for use with TensorFlow is way beyond the scope of this README. There's hardware compatibility to consider, drivers to install... There's a lot. So you're just going to have to figure this part out on your own if you want a GPU. But at least from a Python package perspective, the above requirements file has you covered.
+## Running Tests
 
-**Training Data** You can find our training data here:
-
-https://www.dropbox.com/sh/7aouywaid7xptpq/AAD_-I4hAHrDeiosDAQksnBma?dl=1
-
-Pretty soon, we're going to add this as a TensorFlow DataSet, so you don't need to download this separately like this. It'll also let us version the data a bit better. But for now, just deal with it. There's two things you need from the training data:
-
-1. `images/` folder, containing all the screenshots (resized down to 224x140. We'll have the full-size images up soon)
-2. `labels.csv` that has all the labels
-3. `bishop-fox-pretrained-v1.h5` A pretrained weights file you can use right out of the box without training.
-
-Copy all three into the root of the Eyeballer code tree.
+```
+make test
+```
 
 ## Predicting Labels
 To eyeball some screenshots, just run the "predict" mode:
 
 ```
-eyeballer.py --weights YOUR_WEIGHTS.h5 predict YOUR_FILE.png
+eyeballer --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/
+eyeballer --weights YOUR_WEIGHTS.h5 predict PATH_TO/YOUR_FILES/
 ```
 
-Eyeballer will spit the results back to you in human readable format (a `results.html` file so you can browse it easily) and machine readable format (a `results.csv` file).
-
-
 ## Training
 To train a new model, run:
 ```
-eyeballer.py train
+eyeballer 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.
@@ -85,7 +52,5 @@ This will output a new model file (weights.h5 by default).
 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 --weights YOUR_WEIGHTS.h5 evaluate
 ```
-
-The output will describe the model's accuracy in both recall and precision for each of the program's labels. (Including "none of the above" as a pseudo-label)

eyeballer/__init__.py

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

eyeballer/augmentation.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 from Augmentor.Operations import Operation
-from tensorflow.keras.applications.mobilenet import preprocess_input
+from keras.applications.mobilenet import preprocess_input
 
 
 class EyeballerAugmentation(Operation):

eyeballer/eyeballer.py

@@ -5,7 +5,6 @@ import csv
 
 from eyeballer.model import EyeballModel, DATA_LABELS
 from eyeballer.visualization import HeatMap
-from jinja2 import Template
 
 
 @click.group(invoke_without_command=True)
@@ -38,15 +37,14 @@ 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.option('--threshold', default=.5, type=float, help="Threshold confidence for labeling")
 @click.pass_context
-def predict(ctx, screenshot, heatmap, threshold):
+def predict(ctx, screenshot, heatmap):
     model = EyeballModel(**ctx.obj['model_kwargs'])
     results = model.predict(screenshot)
 
     if heatmap:
         # Generate a heatmap
-        HeatMap(screenshot, model, threshold).generate()
+        HeatMap(screenshot, model).generate()
 
     if not results:
         print("Error: Input file does not exist")
@@ -54,64 +52,24 @@ def predict(ctx, screenshot, heatmap, threshold):
         print(results)
     else:
         with open("results.csv", "w", newline="") as csvfile:
-            fieldnames = ["filename", "custom404", "login", "webapp", "oldlooking"]
+            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]))
+        print(f"{label} Precision Score: {results[label]['precision']:4.2%}")
+        print(f"{label} Recall Score: {results[label]['recall']:4.2%}")
+    print(f"'None of the above' Precision: {results['none_of_the_above_precision']:4.2%}")
+    print(f"'None of the above' Recall: {results['none_of_the_above_recall']:4.2%}")
+    print(f"All or nothing Accuracy: {results['all_or_nothing_accuracy']:4.2%}")
+    print(f"Overall Binary Accuracy: {results['total_binary_accuracy']:4.2%}")
+    print(f"Top 10 worst predictions: {results['top_10_worst'][1]}")
 
 
 @cli.command()

eyeballer/model.py

@@ -1,7 +1,6 @@
 import os
 import random
 import sys
-import progressbar
 
 # Prevent Tkinter Dependency
 import matplotlib
@@ -11,13 +10,34 @@ import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import Augmentor
-import tensorflow as tf
 
-from tensorflow.keras.applications.mobilenet import preprocess_input
+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, hamming_loss
 from eyeballer.augmentation import EyeballerAugmentation
+from keras.callbacks import Callback
 
-DATA_LABELS = ["custom404", "login", "webapp", "oldlooking"]
+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:
@@ -27,11 +47,11 @@ class EyeballModel:
     """
     graphs_directory = "graphs/"
     checkpoint_file = "weights.h5"
-    image_dir = "images/"
+    image_dir = "images_224x224/"
     image_width, image_height = 224, 224
     input_shape = (image_width, image_height, 3)
 
-    def __init__(self, weights_file, print_summary=False, seed=None, quiet=False):
+    def __init__(self, weights_file, print_summary=False, seed=None):
         """Constructor for model class.
 
         Keyword arguments:
@@ -39,40 +59,26 @@ 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
-        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"))
+        # 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)
 
-        self.model.compile(optimizer=tf.keras.optimizers.Adam(0.0005),
-                           loss="binary_crossentropy",
-                           metrics=["accuracy"])
+        # for layer in base_model.layers:
+        #     layer.trainable = False
 
-        if weights_file is not None and os.path.isfile(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:
-            if weights_file is not None:
-                raise FileNotFoundError
-            print("WARN: No model loaded from file. Generating random model")
+        adam = Adam(lr=0.0005)
+        self.model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
 
         if print_summary:
             print(self.model.summary())
 
-        self.quiet = quiet
-        self.seed = seed
-
-    def _init_labels(self):
         # Pull out our labels for use in generators later
         data = pd.read_csv("labels.csv")
         self.training_labels = data.loc[data['evaluation'] == False]  # noqa: E712
@@ -80,14 +86,27 @@ class EyeballModel:
 
         # Shuffle the training labels
         self.random_seed = False
+        self.seed = seed
         if self.seed is None:
             self.random_seed = True
             self.seed = random.randint(0, 999999)
             print("No seed set, ", end='')
-        print("using seed: {}".format(self.seed))
+        print(f"using seed: {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):
+            try:
+                self.model.load_weights(weights_file)
+            except OSError:
+                print(f"ERROR: Unable to open weights file '{weights_file}'")
+                sys.exit(-1)
+            print("Loaded model from file.")
+        else:
+            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)
@@ -112,8 +131,7 @@ class EyeballModel:
         """
         print("Training with seed: " + str(self.seed))
 
-        self._init_labels()
-        data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
+        data_generator = ImageDataGenerator(
             preprocessing_function=self.preprocess_training_function,
             validation_split=0.2,
             samplewise_center=True)
@@ -128,7 +146,7 @@ class EyeballModel:
             subset='training',
             shuffle=True,
             seed=self.seed,
-            class_mode="raw")
+            class_mode="other")
         validation_generator = data_generator.flow_from_dataframe(
             self.training_labels,
             directory=self.image_dir,
@@ -139,26 +157,13 @@ class EyeballModel:
             subset='validation',
             shuffle=False,
             seed=self.seed,
-            class_mode="raw")
+            class_mode="other")
 
         # Model checkpoint - Saves model weights when validation accuracy improves
-        callbacks = [tf.keras.callbacks.ModelCheckpoint(self.checkpoint_file,
-                                                        monitor='val_loss',
-                                                        verbose=1,
-                                                        save_best_only=True,
-                                                        save_weights_only=True,
-                                                        mode='min'),
-                     tf.keras.callbacks.TensorBoard(log_dir='logs',
-                                                    histogram_freq=2,
-                                                    write_graph=True,
-                                                    write_images=False,
-                                                    update_freq='epoch',
-                                                    profile_batch=2,
-                                                    embeddings_freq=0,
-                                                    embeddings_metadata=None),
-                     ]
+        callbacks = [ModelCheckpoint(self.checkpoint_file, monitor='val_loss', verbose=1, save_best_only=True,
+                                     save_weights_only=True, mode='min')]
 
-        history = self.model.fit(
+        history = self.model.fit_generator(
             training_generator,
             steps_per_epoch=len(training_generator.filenames) // batch_size,
             epochs=epochs,
@@ -205,15 +210,15 @@ class EyeballModel:
         result["filename"] = "custom-image"
         result["custom404"] = prediction[0][0]
         result["login"] = prediction[0][1]
-        result["webapp"] = prediction[0][2]
+        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
+    def predict(self, path):
+        """Predict the labels for a single file
 
         Keyword arguments:
-        path -- The path to the file(s) that we'll be evaluating.
+        path -- The path to the file that we'll be evaluating
         """
         # Is this a single file, or a directory?
         screenshots = []
@@ -226,32 +231,22 @@ class EyeballModel:
             raise FileNotFoundError
 
         results = []
-        bar = progressbar.ProgressBar()
-        if self.quiet:
-            bar = progressbar.NullBar()
-        for screenshot in bar(screenshots):
+        for screenshot in 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
+            img = image.load_img(screenshot, target_size=(self.image_width, self.image_height))
+            img = image.img_to_array(img)
+            img = np.expand_dims(img, axis=0)
+            img = preprocess_input(img)
 
             prediction = self.model.predict(img, batch_size=1)
             result = dict()
             result["filename"] = screenshot
             result["custom404"] = prediction[0][0]
             result["login"] = prediction[0][1]
-            result["webapp"] = prediction[0][2]
+            result["homepage"] = prediction[0][2]
             result["oldlooking"] = prediction[0][3]
             results.append(result)
+
         return results
 
     def evaluate(self, threshold=0.5):
@@ -261,8 +256,7 @@ class EyeballModel:
         threshold -- Value between 0->1. The cutoff where the numerical prediction becomes boolean. Default: 0.5
         """
         # Prepare the data
-        self._init_labels()
-        data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
+        data_generator = ImageDataGenerator(
             preprocessing_function=preprocess_input)
         evaluation_generator = data_generator.flow_from_dataframe(
             self.evaluation_labels,
@@ -272,12 +266,12 @@ class EyeballModel:
             target_size=(self.image_width, self.image_height),
             shuffle=False,
             batch_size=1,
-            class_mode="raw")
+            class_mode="other")
         # If a seed was selected, then also evaluate on the validation set for that seed
         if not self.random_seed:
             print("Using validation set...")
             # Data augmentation
-            data_generator = tf.keras.preprocessing.image.ImageDataGenerator(
+            data_generator = ImageDataGenerator(
                 preprocessing_function=self.preprocess_training_function,
                 samplewise_center=True,
                 validation_split=0.2)
@@ -291,24 +285,34 @@ class EyeballModel:
                 subset='validation',
                 shuffle=False,
                 seed=self.seed,
-                class_mode="raw")
+                class_mode="other")
         else:
             print("Using evaluation set...")
 
-        predictions = self.model.predict(
-            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 = self.evaluation_labels[DATA_LABELS].to_numpy()
-        filenames = self.evaluation_labels[["filename"]].to_numpy()
+        ground_truth = evaluation_generator.data
         stats = classification_report(ground_truth, predictions, target_names=DATA_LABELS, output_dict=True)
         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["top_10_best"] = self._top_images(evaluation_generator, predictions, best=True)
+        stats["top_10_worst"] = self._top_images(evaluation_generator, 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
@@ -325,9 +329,6 @@ class EyeballModel:
                 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):
@@ -342,23 +343,20 @@ class EyeballModel:
                 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):
+    def _top_images(self, labels, predictions, top_k=10, best=False):
         """Collect top-k best or top-k worst predicted images
 
         Keyword arguments:
-        ground_truth -- The correct labels
+        labels -- The keras generator that contains the filenames and data 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)
+        true_labels = np.array(labels.data).astype(float)
         predictions = predictions.astype(float)
 
         differences = np.abs(predictions - true_labels).sum(axis=1)
@@ -371,7 +369,7 @@ class EyeballModel:
             indicies = np.flipud(indicies)
 
         for i in indicies[:top_k]:
-            top_file_list.append(filenames[i][0])
+            top_file_list.append(labels.filenames[i])
 
         return indicies[:top_k], top_file_list
 

eyeballer/visualization.py

@@ -1,18 +1,17 @@
 import numpy as np
 import math
-import tensorflow as tf
-
+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):
+    def __init__(self, screenshot_path, model, boxsize=28, step=7):
         self.model = model
 
-        img = tf.keras.preprocessing.image.load_img(screenshot_path, target_size=(self.model.image_width, self.model.image_height))
-        img = tf.keras.preprocessing.image.img_to_array(img)
+        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
@@ -21,7 +20,6 @@ class HeatMap():
         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 """
@@ -32,19 +30,19 @@ class HeatMap():
         for label in DATA_LABELS:
             boxsize = self.boxsize
             # Ignore this label if it didn't predict positively (true label)
-            if results[label] > self.threshold:
+            if results[label] > 0.5:
                 worst_score = 1
-                while worst_score > self.threshold:
+                while worst_score > 0.5:
                     heatmap, worst_score = self._get_heatmap(label, boxsize)
-                    if worst_score > self.threshold:
+                    if worst_score > 0.5:
                         boxsize += 28
-                        print("Didn't get a good image for {}. Trying again with a bigger boxsize: {}".format(label, boxsize))
+                        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 = tf.keras.preprocessing.image.load_img(self.screenshot_path, target_size=(self.model.image_width, self.model.image_width))
+        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
@@ -61,7 +59,7 @@ class HeatMap():
             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))
+        print(f"Heatmap image written to: {output_file}")
 
     def _get_heatmap(self, label, boxsize):
         worst_score = 1

prediction_output_template.html

@@ -1,189 +0,0 @@
-<!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('webapp')"> Web Application</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

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

requirements.txt

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

setup.py

@@ -0,0 +1,33 @@
+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='Bishop Fox',
+    author_email='',
+    description='',
+    entry_points={
+        'console_scripts': [
+            'eyeballer = eyeballer.eyeballer:cli'
+        ]
+    },
+    install_requires=[
+        'Augmentor',
+        'click',
+        'matplotlib',
+        'numpy',
+        'pandas',
+        'pillow',
+        'sklearn',
+        'keras @ git+https://github.com/keras-team/keras',
+    ],
+    extras_require={
+    'tf': ['tensorflow'],
+    'tf_gpu': ['tensorflow-gpu'],
+    },
+    dependency_links=[
+    ],
+)

tests/EvaluateTest.py

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

tests/PredictTest.py

@@ -1,23 +1,15 @@
 import errno
 import os
 import unittest
-import sys
 
 from eyeballer.model import EyeballModel
-from eyeballer.visualization import HeatMap
 
 
 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"
+        weights_file = "latest_weights.h5"
         if not os.path.isfile(weights_file):
-            print("Error: Symlink the latest weights file to " + weights_file)
+            print("Error: Symlink the latest weights file to 'latest_weights.h5'")
             raise FileNotFoundError(
                 errno.ENOENT,
                 os.strerror(errno.ENOENT),
@@ -26,8 +18,7 @@ class PredictTest(unittest.TestCase):
         model_kwargs = {
             "weights_file": weights_file,
             "print_summary": False,
-            "seed": None,
-            "quiet": True
+            "seed": None
         }
         self.model = EyeballModel(**model_kwargs)
 
@@ -35,12 +26,11 @@ class PredictTest(unittest.TestCase):
         model_kwargs = {
             "weights_file": None,
             "print_summary": False,
-            "seed": 12345678,
-            "quiet": True
+            "seed": 12345678
         }
         same_seed_model = EyeballModel(**model_kwargs)
 
-        screenshot = "tests/data/404.png"
+        screenshot = "tests/data/test.png"
 
         results_one = self.model.predict(screenshot)
         results_two = same_seed_model.predict(screenshot)
@@ -51,92 +41,44 @@ class PredictTest(unittest.TestCase):
         model_kwargs = {
             "weights_file": None,
             "print_summary": False,
-            "seed": 12345678,
-            "quiet": True
+            "seed": 12345678
         }
         same_seed_model = EyeballModel(**model_kwargs)
 
-        screenshot = "tests/data/404.png"
+        screenshot = "tests/data/test.png"
 
         results_one = same_seed_model.predict(screenshot)
         results_two = same_seed_model.predict(screenshot)
 
-        self.assertEqual(results_one, results_two)
+        self.assertEquals(results_one, results_two)
+
+    def get_prediction_by_name(self, results):
+        prediction = None
+        max_val = 0.0
+
+        print(results)
+        for k, v in results[0].items():
+            if not isinstance(v, str) and v > max_val:
+                prediction = k
+                max_val = v
+        return prediction
 
     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)
+        screenshot = "tests/data/custom404.png"
+        results = self.model.predict(screenshot)
+        self.assertEqual(self.get_prediction_by_name(results), "custom404")
 
     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)
+        results = self.model.predict(screenshot)
+        self.assertEqual(self.get_prediction_by_name(results), "login")
 
     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)
+        results = self.model.predict(screenshot)
+        self.assertEqual(self.get_prediction_by_name(results), "homepage")
 
     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)
-
-    def test_heatmap(self):
-        screenshot = "tests/data/login.png"
-        HeatMap(screenshot, self.model, 0.5)
-
-        screenshot = "tests/data/404.png"
-        HeatMap(screenshot, self.model, 0.5)
-
-        screenshot = "tests/data/empty.png"
-        HeatMap(screenshot, self.model, 0.5)
+        screenshot = "tests/data/oldlooking.png"
+        results = self.model.predict(screenshot)
+        self.assertEqual(self.get_prediction_by_name(results), "oldlooking")

tests/TrainTest.py

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

tests/__init__.py

@@ -1 +1 @@
-from tests.PredictTest import *  # noqa
+from tests.PredictTest import *

tests/data/invalid.png

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