Added a GPU check for NVIDIA in Makefile

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

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

@@ -0,0 +1,56 @@
+# 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.
+
+## Installation
+
+To install in a python venv (Recommended):
+```
+python3 -m venv venv
+source venv/bin/activate
+make init
+```
+
+To install for Conda (automatically creates conda venv):
+
+```
+make conda
+```
+
+## Running Tests
+
+```
+make test
+```
+
+## Predicting Labels
+To eyeball some screenshots, just run the "predict" mode:
+
+```
+eyeballer --weights YOUR_WEIGHTS.h5 predict YOUR_FILE.png
+```
+
+Or for a whole directory of files:
+
+```
+eyeballer --weights YOUR_WEIGHTS.h5 predict PATH_TO/YOUR_FILES/
+```
+
+## Training
+To train a new model, run:
+```
+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.
+
+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 --weights YOUR_WEIGHTS.h5 evaluate
+```

eyeballer/__init__.py

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

eyeballer/conv_filter_visualization.py

@@ -1,263 +0,0 @@
-"""
-#from https://raw.githubusercontent.com/keras-team/keras/master/examples/conv_filter_visualization.py
-
-#Visualization of the filters of VGG16, via gradient ascent in input space.
-
-This script can run on CPU in a few minutes.
-
-Results example: ![Visualization](http://i.imgur.com/4nj4KjN.jpg)
-"""
-from __future__ import print_function
-
-import time
-import numpy as np
-from PIL import Image as pil_image
-from keras.preprocessing.image import save_img
-from keras import layers
-from keras.applications import vgg16
-from keras import backend as K
-
-
-def normalize(x):
-    """utility function to normalize a tensor.
-
-    # Arguments
-        x: An input tensor.
-
-    # Returns
-        The normalized input tensor.
-    """
-    return x / (K.sqrt(K.mean(K.square(x))) + K.epsilon())
-
-
-def deprocess_image(x):
-    """utility function to convert a float array into a valid uint8 image.
-
-    # Arguments
-        x: A numpy-array representing the generated image.
-
-    # Returns
-        A processed numpy-array, which could be used in e.g. imshow.
-    """
-    # normalize tensor: center on 0., ensure std is 0.25
-    x -= x.mean()
-    x /= (x.std() + K.epsilon())
-    x *= 0.25
-
-    # clip to [0, 1]
-    x += 0.5
-    x = np.clip(x, 0, 1)
-
-    # convert to RGB array
-    x *= 255
-    if K.image_data_format() == 'channels_first':
-        x = x.transpose((1, 2, 0))
-    x = np.clip(x, 0, 255).astype('uint8')
-    return x
-
-
-def process_image(x, former):
-    """utility function to convert a valid uint8 image back into a float array.
-       Reverses `deprocess_image`.
-
-    # Arguments
-        x: A numpy-array, which could be used in e.g. imshow.
-        former: The former numpy-array.
-                Need to determine the former mean and variance.
-
-    # Returns
-        A processed numpy-array representing the generated image.
-    """
-    if K.image_data_format() == 'channels_first':
-        x = x.transpose((2, 0, 1))
-    return (x / 255 - 0.5) * 4 * former.std() + former.mean()
-
-
-def visualize_layer(model,
-                    layer_name,
-                    step=1.,
-                    epochs=15,
-                    upscaling_steps=9,
-                    upscaling_factor=1.2,
-                    output_dim=(412, 412),
-                    filter_range=(0, None)):
-    """Visualizes the most relevant filters of one conv-layer in a certain model.
-
-    # Arguments
-        model: The model containing layer_name.
-        layer_name: The name of the layer to be visualized.
-                    Has to be a part of model.
-        step: step size for gradient ascent.
-        epochs: Number of iterations for gradient ascent.
-        upscaling_steps: Number of upscaling steps.
-                         Starting image is in this case (80, 80).
-        upscaling_factor: Factor to which to slowly upgrade
-                          the image towards output_dim.
-        output_dim: [img_width, img_height] The output image dimensions.
-        filter_range: Tupel[lower, upper]
-                      Determines the to be computed filter numbers.
-                      If the second value is `None`,
-                      the last filter will be inferred as the upper boundary.
-    """
-
-    def _generate_filter_image(input_img,
-                               layer_output,
-                               filter_index):
-        """Generates image for one particular filter.
-
-        # Arguments
-            input_img: The input-image Tensor.
-            layer_output: The output-image Tensor.
-            filter_index: The to be processed filter number.
-                          Assumed to be valid.
-
-        #Returns
-            Either None if no image could be generated.
-            or a tuple of the image (array) itself and the last loss.
-        """
-        s_time = time.time()
-
-        # we build a loss function that maximizes the activation
-        # of the nth filter of the layer considered
-        if K.image_data_format() == 'channels_first':
-            loss = K.mean(layer_output[:, filter_index, :, :])
-        else:
-            loss = K.mean(layer_output[:, :, :, filter_index])
-
-        # we compute the gradient of the input picture wrt this loss
-        grads = K.gradients(loss, input_img)[0]
-
-        # normalization trick: we normalize the gradient
-        grads = normalize(grads)
-
-        # this function returns the loss and grads given the input picture
-        iterate = K.function([input_img], [loss, grads])
-
-        # we start from a gray image with some random noise
-        intermediate_dim = tuple(
-            int(x / (upscaling_factor ** upscaling_steps)) for x in output_dim)
-        if K.image_data_format() == 'channels_first':
-            input_img_data = np.random.random(
-                (1, 3, intermediate_dim[0], intermediate_dim[1]))
-        else:
-            input_img_data = np.random.random(
-                (1, intermediate_dim[0], intermediate_dim[1], 3))
-        input_img_data = (input_img_data - 0.5) * 20 + 128
-
-        # Slowly upscaling towards the original size prevents
-        # a dominating high-frequency of the to visualized structure
-        # as it would occur if we directly compute the 412d-image.
-        # Behaves as a better starting point for each following dimension
-        # and therefore avoids poor local minima
-        for up in reversed(range(upscaling_steps)):
-            # we run gradient ascent for e.g. 20 steps
-            for _ in range(epochs):
-                loss_value, grads_value = iterate([input_img_data])
-                input_img_data += grads_value * step
-
-                # some filters get stuck to 0, we can skip them
-                if loss_value <= K.epsilon():
-                    return None
-
-            # Calulate upscaled dimension
-            intermediate_dim = tuple(
-                int(x / (upscaling_factor ** up)) for x in output_dim)
-            # Upscale
-            img = deprocess_image(input_img_data[0])
-            img = np.array(pil_image.fromarray(img).resize(intermediate_dim,
-                                                           pil_image.BICUBIC))
-            input_img_data = [process_image(img, input_img_data[0])]
-
-        # decode the resulting input image
-        img = deprocess_image(input_img_data[0])
-        e_time = time.time()
-        print('Costs of filter {:3}: {:5.0f} ( {:4.2f}s )'.format(filter_index,
-                                                                  loss_value,
-                                                                  e_time - s_time))
-        return img, loss_value
-
-    def _draw_filters(filters, n=None):
-        """Draw the best filters in a nxn grid.
-
-        # Arguments
-            filters: A List of generated images and their corresponding losses
-                     for each processed filter.
-            n: dimension of the grid.
-               If none, the largest possible square will be used
-        """
-        if n is None:
-            n = int(np.floor(np.sqrt(len(filters))))
-
-        # the filters that have the highest loss are assumed to be better-looking.
-        # we will only keep the top n*n filters.
-        filters.sort(key=lambda x: x[1], reverse=True)
-        filters = filters[:n * n]
-
-        # build a black picture with enough space for
-        # e.g. our 8 x 8 filters of size 412 x 412, with a 5px margin in between
-        MARGIN = 5
-        width = n * output_dim[0] + (n - 1) * MARGIN
-        height = n * output_dim[1] + (n - 1) * MARGIN
-        stitched_filters = np.zeros((width, height, 3), dtype='uint8')
-
-        # fill the picture with our saved filters
-        for i in range(n):
-            for j in range(n):
-                img, _ = filters[i * n + j]
-                width_margin = (output_dim[0] + MARGIN) * i
-                height_margin = (output_dim[1] + MARGIN) * j
-                stitched_filters[
-                    width_margin: width_margin + output_dim[0],
-                    height_margin: height_margin + output_dim[1], :] = img
-
-        # save the result to disk
-        save_img('viz_{0:}_{1:}x{1:}.png'.format(layer_name, n), stitched_filters)
-
-    # this is the placeholder for the input images
-    assert len(model.inputs) == 1
-    input_img = model.inputs[0]
-
-    # get the symbolic outputs of each "key" layer (we gave them unique names).
-    layer_dict = dict([(layer.name, layer) for layer in model.layers[1:]])
-
-    output_layer = layer_dict[layer_name]
-    assert isinstance(output_layer, layers.Conv2D)
-
-    # Compute to be processed filter range
-    import pdb
-    pdb.set_trace()
-    filter_lower = filter_range[0]
-    filter_upper = (filter_range[1]
-                    if filter_range[1] is not None
-                    else len(output_layer.get_weights()[0][0][0][0]))
-                    # BF: Using size instead of [0] to account for n-dim array
-    assert(filter_lower >= 0
-           and filter_upper <= len(output_layer.get_weights()[0][0][0][0])
-           and filter_upper > filter_lower)
-    print('Compute filters {:} to {:}'.format(filter_lower, filter_upper))
-
-    # iterate through each filter and generate its corresponding image
-    processed_filters = []
-    for f in range(filter_lower, filter_upper):
-        img_loss = _generate_filter_image(input_img, output_layer.output, f)
-
-        if img_loss is not None:
-            processed_filters.append(img_loss)
-
-    print('{} filter processed.'.format(len(processed_filters)))
-    # Finally draw and store the best filters to disk
-    _draw_filters(processed_filters)
-
-
-if __name__ == '__main__':
-    # the name of the layer we want to visualize
-    # (see model definition at keras/applications/vgg16.py)
-    LAYER_NAME = 'block5_conv1'
-
-    # build the VGG16 network with ImageNet weights
-    vgg = vgg16.VGG16(weights='imagenet', include_top=False)
-    print('Model loaded.')
-    vgg.summary()
-
-    # example function call
-    visualize_layer(vgg, LAYER_NAME)

eyeballer/deep_dream.py

@@ -1,182 +0,0 @@
-'''
-#Deep Dreaming in Keras.
-
-Run the script with:
-```python
-python deep_dream.py path_to_your_base_image.jpg prefix_for_results
-```
-e.g.:
-```python
-python deep_dream.py img/mypic.jpg results/dream
-```
-'''
-from __future__ import print_function
-
-from keras.preprocessing.image import load_img, save_img, img_to_array
-import numpy as np
-import scipy
-import argparse
-
-from keras.applications import inception_v3
-from keras import backend as K
-
-fetch_loss_and_grads = None
-# These are the names of the layers
-# for which we try to maximize activation,
-# as well as their weight in the final loss
-# we try to maximize.
-# You can tweak these setting to obtain new visual effects.
-settings = {
-    'features': {
-        'conv1': 1.,
-        #'mixed2': 0.2,
-        #'mixed3': 0.5,
-        #'mixed4': 2.,
-        #'mixed5': 1.5,
-    },
-}
-
-
-def preprocess_image(image_path):
-    # Util function to open, resize and format pictures
-    # into appropriate tensors.
-    img = load_img(image_path)
-    img = img_to_array(img)
-    img = np.expand_dims(img, axis=0)
-    img = inception_v3.preprocess_input(img)
-    return img
-
-
-def deprocess_image(x):
-    # Util function to convert a tensor into a valid image.
-    if K.image_data_format() == 'channels_first':
-        x = x.reshape((3, x.shape[2], x.shape[3]))
-        x = x.transpose((1, 2, 0))
-    else:
-        x = x.reshape((x.shape[1], x.shape[2], 3))
-    x /= 2.
-    x += 0.5
-    x *= 255.
-    x = np.clip(x, 0, 255).astype('uint8')
-    return x
-
-
-
-def eval_loss_and_grads(x):
-    outs = fetch_loss_and_grads([x])
-    loss_value = outs[0]
-    grad_values = outs[1]
-    return loss_value, grad_values
-
-
-def resize_img(img, size):
-    img = np.copy(img)
-    if K.image_data_format() == 'channels_first':
-        factors = (1, 1,
-                   float(size[0]) / img.shape[2],
-                   float(size[1]) / img.shape[3])
-    else:
-        factors = (1,
-                   float(size[0]) / img.shape[1],
-                   float(size[1]) / img.shape[2],
-                   1)
-    return scipy.ndimage.zoom(img, factors, order=1)
-
-
-def gradient_ascent(x, iterations, step, max_loss=None):
-    for i in range(iterations):
-        loss_value, grad_values = eval_loss_and_grads(x)
-        if max_loss is not None and loss_value > max_loss:
-            break
-        print('..Loss value at', i, ':', loss_value)
-        x += step * grad_values
-    return x
-
-
-"""Process:
-
-- Load the original image.
-- Define a number of processing scales (i.e. image shapes),
-    from smallest to largest.
-- Resize the original image to the smallest scale.
-- For every scale, starting with the smallest (i.e. current one):
-    - Run gradient ascent
-    - Upscale image to the next scale
-    - Reinject the detail that was lost at upscaling time
-- Stop when we are back to the original size.
-
-To obtain the detail lost during upscaling, we simply
-take the original image, shrink it down, upscale it,
-and compare the result to the (resized) original image.
-"""
-def process(model, base_image_path, result_prefix):
-    global settings
-    global fetch_loss_and_grads
-    K.set_learning_phase(0)
-
-    dream = model.input
-    print('Model loaded.')
-
-    # Get the symbolic outputs of each "key" layer (we gave them unique names).
-    layer_dict = dict([(layer.name, layer) for layer in model.layers])
-
-    # Define the loss.
-    loss = K.variable(0.)
-    for layer_name in settings['features']:
-        # Add the L2 norm of the features of a layer to the loss.
-        if layer_name not in layer_dict:
-            raise ValueError('Layer ' + layer_name + ' not found in model.')
-        coeff = settings['features'][layer_name]
-        x = layer_dict[layer_name].output
-        # We avoid border artifacts by only involving non-border pixels in the loss.
-        scaling = K.prod(K.cast(K.shape(x), 'float32'))
-        if K.image_data_format() == 'channels_first':
-            loss += coeff * K.sum(K.square(x[:, :, 2: -2, 2: -2])) / scaling
-        else:
-            loss += coeff * K.sum(K.square(x[:, 2: -2, 2: -2, :])) / scaling
-
-    # Compute the gradients of the dream wrt the loss.
-    grads = K.gradients(loss, dream)[0]
-    # Normalize gradients.
-    grads /= K.maximum(K.mean(K.abs(grads)), K.epsilon())
-
-    # Set up function to retrieve the value
-    # of the loss and gradients given an input image.
-    outputs = [loss, grads]
-    fetch_loss_and_grads = K.function([dream], outputs)
-
-    # Playing with these hyperparameters will also allow you to achieve new effects
-    step = 0.01  # Gradient ascent step size
-    num_octave = 3  # Number of scales at which to run gradient ascent
-    octave_scale = 1.4  # Size ratio between scales
-    iterations = 20  # Number of ascent steps per scale
-    max_loss = 10.
-
-    img = preprocess_image(base_image_path)
-    if K.image_data_format() == 'channels_first':
-        original_shape = img.shape[2:]
-    else:
-        original_shape = img.shape[1:3]
-    successive_shapes = [original_shape]
-    for i in range(1, num_octave):
-        shape = tuple([int(dim / (octave_scale ** i)) for dim in original_shape])
-        successive_shapes.append(shape)
-    successive_shapes = successive_shapes[::-1]
-    original_img = np.copy(img)
-    shrunk_original_img = resize_img(img, successive_shapes[0])
-
-    for shape in successive_shapes:
-        print('Processing image shape', shape)
-        img = resize_img(img, shape)
-        img = gradient_ascent(img,
-                            iterations=iterations,
-                            step=step,
-                            max_loss=max_loss)
-        upscaled_shrunk_original_img = resize_img(shrunk_original_img, shape)
-        same_size_original = resize_img(original_img, shape)
-        lost_detail = same_size_original - upscaled_shrunk_original_img
-
-        img += lost_detail
-        shrunk_original_img = resize_img(original_img, shape)
-
-    save_img(result_prefix + '.png', deprocess_image(np.copy(img)))

eyeballer/eyeballer.py

@@ -1,11 +1,10 @@
 #!/usr/bin/env python3
+
 import click
 import csv
-from conv_filter_visualization import visualize_layer
-from keras import layers
-import deep_dream
 
-from model import EyeballModel, DATA_LABELS
+from eyeballer.model import EyeballModel, DATA_LABELS
+from eyeballer.visualization import HeatMap
 
 
 @click.group(invoke_without_command=True)
@@ -37,10 +36,16 @@ def train(ctx, graphs, batchsize, epochs):
 
 @cli.command()
 @click.argument('screenshot')
+@click.option('--heatmap', default=False, is_flag=True, help="Create a heatmap graphfor the prediction")
 @click.pass_context
-def predict(ctx, screenshot):
+def predict(ctx, screenshot, heatmap):
     model = EyeballModel(**ctx.obj['model_kwargs'])
     results = model.predict(screenshot)
+
+    if heatmap:
+        # Generate a heatmap
+        HeatMap(screenshot, model).generate()
+
     if not results:
         print("Error: Input file does not exist")
     if len(results) == 1:
@@ -60,8 +65,10 @@ def pretty_print_evaluation(results):
     for label in DATA_LABELS:
         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"All or nothing accuracy: {results['all_or_nothing_accuracy']:4.2%}")
     print(f"Top 10 worst predictions: {results['top_10_worst'][1]}")
 
 
@@ -74,29 +81,5 @@ def evaluate(ctx, threshold):
     pretty_print_evaluation(results)
 
 
-@cli.command()
-@click.argument('layer', default=None, required=False)
-@click.pass_context
-def visualize(ctx, layer):
-    model = EyeballModel(**ctx.obj['model_kwargs']).model
-    if not layer:
-        layer_dict = dict([(layer.name, layer) for layer in model.layers[1:]])
-        print("Layers:")
-        for k in layer_dict.keys():
-            if isinstance(layer_dict[k], layers.Conv2D):
-                print("\t" + k)
-    else:
-        visualize_layer(model, layer)
-
-@cli.command()
-@click.option('--base-image-path', type=click.Path(),
-                    help='Path to the image to transform.')
-@click.option('--result-prefix', type=click.Path(),
-                    help='Prefix for the saved results.')
-@click.pass_context
-def deepdream(ctx, base_image_path, result_prefix):
-    model = EyeballModel(**ctx.obj['model_kwargs']).model
-    deep_dream.process(model, base_image_path, result_prefix)
-
 if __name__ == '__main__':
     cli()

eyeballer/model.py

@@ -1,5 +1,6 @@
 import os
 import random
+import sys
 
 # Prevent Tkinter Dependency
 import matplotlib
@@ -21,7 +22,7 @@ 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 augmentation import EyeballerAugmentation
+from eyeballer.augmentation import EyeballerAugmentation
 from keras.callbacks import Callback
 
 DATA_LABELS = ["custom404", "login", "homepage", "oldlooking"]
@@ -50,7 +51,7 @@ class EyeballModel:
     image_width, image_height = 224, 224
     input_shape = (image_width, image_height, 3)
 
-    def __init__(self, print_summary=False, weights_file="weights.h5", seed=None):
+    def __init__(self, weights_file, print_summary=False, seed=None):
         """Constructor for model class.
 
         Keyword arguments:
@@ -95,10 +96,16 @@ class EyeballModel:
         self.training_labels = self.training_labels.sample(frac=1)
 
         if weights_file is not None and os.path.isfile(weights_file):
-            self.model.load_weights(weights_file)
+            try:
+                self.model.load_weights(weights_file)
+            except OSError:
+                print(f"ERROR: Unable to open weights file '{weights_file}'")
+                sys.exit(-1)
             print("Loaded model from file.")
         else:
-            print("No model loaded from file")
+            if weights_file is not None:
+                raise FileNotFoundError
+            print("WARN: No model loaded from file. Generating random model")
 
         # Data augmentation
         augmentor = Augmentor.Pipeline()
@@ -189,6 +196,24 @@ class EyeballModel:
             plt.legend(['Train', 'Validation'], loc='upper left')
             plt.savefig(self.graphs_directory + "loss.png")
 
+    def predict_on_array(self, image):
+        """Predict the labels for a single screenshot
+
+        Keyword arguments:
+        image -- The numpy array of the image to classify
+        """
+        img = np.expand_dims(image, axis=0)
+        img = preprocess_input(img)
+
+        prediction = self.model.predict(img, batch_size=1)
+        result = dict()
+        result["filename"] = "custom-image"
+        result["custom404"] = prediction[0][0]
+        result["login"] = prediction[0][1]
+        result["homepage"] = prediction[0][2]
+        result["oldlooking"] = prediction[0][3]
+        return result
+
     def predict(self, path):
         """Predict the labels for a single file
 
@@ -203,7 +228,7 @@ class EyeballModel:
             screenshots = os.listdir(path)
             screenshots = [os.path.join(path, s) for s in screenshots]
         else:
-            return None
+            raise FileNotFoundError
 
         results = []
         for screenshot in screenshots:
@@ -288,8 +313,38 @@ class EyeballModel:
         stats["all_or_nothing_accuracy"] = accuracy_score(ground_truth, predictions)
         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
 
+    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
+        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
+        return correct_count / total_count
+
     def _top_images(self, labels, predictions, top_k=10, best=False):
         """Collect top-k best or top-k worst predicted images
 

eyeballer/visualization.py

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

requirements.txt

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

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

@@ -0,0 +1,84 @@
+import errno
+import os
+import unittest
+
+from eyeballer.model import EyeballModel
+
+
+class PredictTest(unittest.TestCase):
+    def setUp(self):
+        weights_file = "latest_weights.h5"
+        if not os.path.isfile(weights_file):
+            print("Error: Symlink the latest weights file to 'latest_weights.h5'")
+            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/test.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/test.png"
+
+        results_one = same_seed_model.predict(screenshot)
+        results_two = same_seed_model.predict(screenshot)
+
+        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/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)
+        self.assertEqual(self.get_prediction_by_name(results), "login")
+
+    def test_predict_homepage(self):
+        screenshot = "tests/data/homepage.png"
+        results = self.model.predict(screenshot)
+        self.assertEqual(self.get_prediction_by_name(results), "homepage")
+
+    def test_predict_oldlooking(self):
+        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

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