Adding DataViz Driver

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

.gitignore

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

eyeballer/__init__.py

@@ -1 +1 @@
-from eyeballer.eyeballer import EyeballModel
+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/conv_filter_visualization.py

@@ -0,0 +1,263 @@
+"""
+#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

@@ -0,0 +1,182 @@
+'''
+#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,274 +1,102 @@
-#!/usr/bin/python3
+#!/usr/bin/env python3
+import click
+import csv
+from conv_filter_visualization import visualize_layer
+from keras import layers
+import deep_dream
 
-import argparse
-import os
-import random
-import time
-
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-
-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 precision_score, recall_score
+from model import EyeballModel, DATA_LABELS
 
 
-class EyeballModel:
-    """The primary model class of Eyeballer.
+@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}
 
-    Contains high-level functions for training, evaluating, and predicting.
-    """
-    checkpoint_file = "weights-snapshot.h5"
-    image_dir = "images_224x224/"
-    image_width, image_height = 224, 224
-    input_shape = (image_width, image_height, 3)
+    #  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
 
-    def __init__(self, print_summary=False, weights_file="weights.h5", seed=None):
-        """Constructor for model class.
 
-        Keyword arguments:
-        print_summary -- Whether or not to print to stdout the keras model summary, containing a detailed description of every model layer
-        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)
+@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)
 
-        # for layer in base_model.layers:
-        #    layer.trainable = False
 
-        adam = Adam(lr=0.0005)
-        self.model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
+@cli.command()
+@click.argument('screenshot')
+@click.pass_context
+def predict(ctx, screenshot):
+    model = EyeballModel(**ctx.obj['model_kwargs'])
+    results = model.predict(screenshot)
+    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")
 
-        if print_summary:
-            print(self.model.summary())
 
-        # Pull out our labels for use in generators later
-        data = pd.read_csv("labels.csv")
-        self.training_labels = data.loc[data['evaluation'] == False]
-        self.evaluation_labels = data.loc[data['evaluation'] == True]
+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} Recall 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%}")
+    print(f"Top 10 worst predictions: {results['top_10_worst'][1]}")
 
-        # Shuffle the training labels
-        self.seed = seed
-        if self.seed is None:
-            self.seed = 0
-            print("No seed set, ", end='')
-        print(f"using seed: {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)
-            print("Loaded model from file.")
-        else:
-            print("No model loaded from file")
+@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)
 
-    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.
 
-        Keyword arguments:
-        epochs -- The number of epochs to train for. (An epoch is one pass-through of the dataset)
-        batch_size -- How many images to batch together when training. Generally speaking, the higher the better, until you run out of memory.
-        print_graphs --- Whether or not to create accuracy and loss graphs. If true, they'll be written to accuracy.png and loss.png
-        """
-        print("Training with seed: " + str(self.seed))
+@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)
 
-        # Data augmentation
-        data_generator = ImageDataGenerator(
-            preprocessing_function=preprocess_input,
-            shear_range=0.0,
-            zoom_range=0.2,
-            samplewise_center=True,
-            validation_split=0.2,
-            horizontal_flip=False)
-
-        training_generator = data_generator.flow_from_dataframe(
-            self.training_labels,
-            directory=self.image_dir,
-            x_col="filename",
-            y_col=["custom404", "login", "homepage"],
-            target_size=(self.image_width, self.image_height),
-            batch_size=batch_size,
-            subset='training',
-            shuffle=True,
-            seed=self.seed,
-            class_mode="other")
-        validation_generator = data_generator.flow_from_dataframe(
-            self.training_labels,
-            directory=self.image_dir,
-            x_col="filename",
-            y_col=["custom404", "login", "homepage"],
-            target_size=(self.image_width, self.image_height),
-            batch_size=batch_size,
-            subset='validation',
-            shuffle=False,
-            seed=self.seed,
-            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')]
-
-        history = self.model.fit_generator(
-            training_generator,
-            steps_per_epoch=len(training_generator.filenames) // batch_size,
-            epochs=epochs,
-            validation_data=validation_generator,
-            validation_steps=len(validation_generator.filenames) // batch_size,
-            callbacks=callbacks,
-            verbose=1)
-
-        if print_graphs:
-            # Plot training & validation accuracy values
-            plt.plot(history.history['acc'])
-            plt.plot(history.history['val_acc'])
-            plt.title('Model accuracy')
-            plt.ylabel('Accuracy')
-            plt.xlabel('Epoch')
-            plt.legend(['Train', 'Validation'], loc='upper left')
-            plt.savefig("accuracy.png")
-            plt.clf()
-            plt.cla()
-            plt.close()
-
-            # Plot training & validation loss values
-            plt.plot(history.history['loss'])
-            plt.plot(history.history['val_loss'])
-            plt.title('Model loss')
-            plt.ylabel('Loss')
-            plt.xlabel('Epoch')
-            plt.legend(['Train', 'Validation'], loc='upper left')
-            plt.savefig("loss.png")
-
-    def predict(self, filename):
-        """Predict the labels for a single file
-
-        Keyword arguments:
-        filename -- The path to the file that we'll be evaluating
-        """
-        # 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 = 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)
-
-    def evaluate(self, threshold=0.5):
-        """Evaluate performance against the persistent evaluation data set
-
-        Keyword arguments:
-        threshold -- Value between 0->1. The cutoff where the numerical prediction becomes boolean. Default: 0.5
-        """
-        # Prepare the data
-        data_generator = ImageDataGenerator(
-            preprocessing_function=preprocess_input)
-        evaluation_generator = data_generator.flow_from_dataframe(
-            self.evaluation_labels,
-            directory=self.image_dir,
-            x_col="filename",
-            y_col=["custom404", "login", "homepage"],
-            target_size=(self.image_width, self.image_height),
-            shuffle=False,
-            batch_size=1,
-            class_mode="other")
-        # If a seed was selected, then also evaluate on the validation set for that seed
-        if self.seed is not None:
-            print("Using validation set...")
-            # Data augmentation
-            data_generator = ImageDataGenerator(
-                preprocessing_function=preprocess_input,
-                shear_range=0.0,
-                zoom_range=0.2,
-                samplewise_center=True,
-                validation_split=0.2,
-                horizontal_flip=False)
-            evaluation_generator = data_generator.flow_from_dataframe(
-                self.training_labels,
-                directory=self.image_dir,
-                x_col="filename",
-                y_col=["custom404", "login", "homepage"],
-                target_size=(self.image_width, self.image_height),
-                batch_size=1,
-                subset='validation',
-                shuffle=False,
-                seed=self.seed,
-                class_mode="other")
-        else:
-            print("Using evaluation set...")
-
-        predictions = self.model.predict_generator(evaluation_generator, verbose=1, steps=len(evaluation_generator))
-        predictions = predictions > threshold
-
-        all_or_nothing_success = 0
-
-        correct_per_category_counts = [0, 0, 0]
-        total_correct_count = 0
-
-        labels = evaluation_generator.data.tolist()
-
-        # Zip them together, so we can evaluate easier
-        for row in zip(labels, predictions.tolist()):
-            # All or nothing count
-            if row[0] == row[1]:
-                all_or_nothing_success += 1
-
-            # Per-category count
-            for i in range(len(row[0])):
-                if row[0][i] == row[1][i]:
-                    correct_per_category_counts[i] += 1
-                    total_correct_count += 1
-
-        custom404_labels = np.reshape(np.array(labels)[:, 0:1], len(labels)).tolist()
-        custom404_predicitons = np.reshape(np.array(predictions)[:, 0:1], len(predictions)).tolist()
-
-        login_labels = np.reshape(np.array(labels)[:, 1:2], len(labels)).tolist()
-        login_predicitons = np.reshape(np.array(predictions)[:, 1:2], len(predictions)).tolist()
-
-        homepage_labels = np.reshape(np.array(labels)[:, 2:3], len(labels)).tolist()
-        homepage_predicitons = np.reshape(np.array(predictions)[:, 2:3], len(predictions)).tolist()
-
-        custom404_precision = precision_score(custom404_labels, custom404_predicitons)
-        custom404_recall = recall_score(custom404_labels, custom404_predicitons)
-        print("Custom404 Precision Score: " + str(round(custom404_precision * 100, 2)) + "%")
-        print("Custom404 Recall Score: " + str(round(custom404_recall * 100, 2)) + "%")
-
-        login_precision = precision_score(login_labels, login_predicitons)
-        login_recall = recall_score(login_labels, login_predicitons)
-        print("Login Precision Score: " + str(round(login_precision * 100, 2)) + "%")
-        print("Login Recall Score: " + str(round(login_recall * 100, 2)) + "%")
-
-        homepage_precision = precision_score(homepage_labels, homepage_predicitons)
-        homepage_recall = recall_score(homepage_labels, homepage_predicitons)
-        print("Homepage Precision Score: " + str(round(homepage_precision * 100, 2)) + "%")
-        print("Homepage Recall Score: " + str(round(homepage_recall * 100, 2)) + "%")
-
-        print("Overall Binary Accuracy: " + str(
-            round((total_correct_count * 100) / (len(evaluation_generator) * 3), 2)) + "%")
-        print("All or nothing accuracy: " + str(
-            round((all_or_nothing_success / len(evaluation_generator)) * 100, 2)) + "%")
+@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/main.py

@@ -1,58 +0,0 @@
-#!/env/python
-import click
-import random
-from eyeballer import EyeballModel
-
-
-
-
-@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, help="RNG seed for data shuffling and transformations, defaults to random value")
-@click.pass_context
-def cli(ctx, weights, summary, seed):
-    if not seed:
-        seed = random.randint(0, 999999)
-
-    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)
-
-
-@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'])
-    model.evaluate(threshold)
-
-
-if __name__ == '__main__':
-    cli()
-

eyeballer/model.py

@@ -0,0 +1,340 @@
+import os
+import random
+
+# 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
+
+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 augmentation import EyeballerAugmentation
+from keras.callbacks import Callback
+
+DATA_LABELS = ["custom404", "login", "homepage", "oldlooking"]
+
+
+class PredictionManager(Callback):
+    def __init__(self):
+        super(PredictionManager, self).__init__()
+        self.predictions = []
+
+    def getPredictions(self):
+        return np.array(self.predictions)
+
+    def on_predict_batch_end(self, batch, logs=None):
+        self.predictions.append(logs['outputs'][0].tolist()[0])
+
+
+class EyeballModel:
+    """The primary model class of Eyeballer.
+
+    Contains high-level functions for training, evaluating, and predicting.
+    """
+    graphs_directory = "graphs/"
+    checkpoint_file = "weights.h5"
+    image_dir = "images_224x224/"
+    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):
+        """Constructor for model class.
+
+        Keyword arguments:
+        print_summary -- Whether or not to print to stdout the keras model summary, containing a detailed description of every model layer
+        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(len(DATA_LABELS), activation="sigmoid", name="OutputLayer")(x)
+        self.model = Model(inputs=base_model.input, outputs=output_layer)
+
+        # for layer in base_model.layers:
+        #     layer.trainable = False
+
+        adam = Adam(lr=0.0005)
+        self.model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
+
+        if print_summary:
+            print(self.model.summary())
+
+        # 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
+        self.evaluation_labels = data.loc[data['evaluation'] == True]  # noqa: E712
+
+        # 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(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):
+            self.model.load_weights(weights_file)
+            print("Loaded model from file.")
+        else:
+            print("No model loaded from file")
+
+        # 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.
+
+        Keyword arguments:
+        epochs -- The number of epochs to train for. (An epoch is one pass-through of the dataset)
+        batch_size -- How many images to batch together when training. Generally speaking, the higher the better, until you run out of memory.
+        print_graphs --- Whether or not to create accuracy and loss graphs. If true, they'll be written to accuracy.png and loss.png
+        """
+        print("Training with seed: " + str(self.seed))
+
+        data_generator = ImageDataGenerator(
+            preprocessing_function=self.preprocess_training_function,
+            validation_split=0.2,
+            samplewise_center=True)
+
+        training_generator = data_generator.flow_from_dataframe(
+            self.training_labels,
+            directory=self.image_dir,
+            x_col="filename",
+            y_col=DATA_LABELS,
+            target_size=(self.image_width, self.image_height),
+            batch_size=batch_size,
+            subset='training',
+            shuffle=True,
+            seed=self.seed,
+            class_mode="other")
+        validation_generator = data_generator.flow_from_dataframe(
+            self.training_labels,
+            directory=self.image_dir,
+            x_col="filename",
+            y_col=DATA_LABELS,
+            target_size=(self.image_width, self.image_height),
+            batch_size=batch_size,
+            subset='validation',
+            shuffle=False,
+            seed=self.seed,
+            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')]
+
+        history = self.model.fit_generator(
+            training_generator,
+            steps_per_epoch=len(training_generator.filenames) // batch_size,
+            epochs=epochs,
+            validation_data=validation_generator,
+            validation_steps=len(validation_generator.filenames) // batch_size,
+            callbacks=callbacks,
+            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'])
+            plt.title('Model accuracy')
+            plt.ylabel('Accuracy')
+            plt.xlabel('Epoch')
+            plt.legend(['Train', 'Validation'], loc='upper left')
+            plt.savefig(self.graphs_directory + "accuracy.png")
+            plt.clf()
+            plt.cla()
+            plt.close()
+
+            # Plot training & validation loss values
+            plt.plot(history.history['loss'])
+            plt.plot(history.history['val_loss'])
+            plt.title('Model loss')
+            plt.ylabel('Loss')
+            plt.xlabel('Epoch')
+            plt.legend(['Train', 'Validation'], loc='upper left')
+            plt.savefig(self.graphs_directory + "loss.png")
+
+    def predict(self, path):
+        """Predict the labels for a single file
+
+        Keyword arguments:
+        path -- The path to the file 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:
+            return None
+
+        results = []
+        for screenshot in screenshots:
+            # Load the image into memory
+            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["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
+
+        Keyword arguments:
+        threshold -- Value between 0->1. The cutoff where the numerical prediction becomes boolean. Default: 0.5
+        """
+        # Prepare the data
+        data_generator = ImageDataGenerator(
+            preprocessing_function=preprocess_input)
+        evaluation_generator = data_generator.flow_from_dataframe(
+            self.evaluation_labels,
+            directory=self.image_dir,
+            x_col="filename",
+            y_col=DATA_LABELS,
+            target_size=(self.image_width, self.image_height),
+            shuffle=False,
+            batch_size=1,
+            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 = ImageDataGenerator(
+                preprocessing_function=self.preprocess_training_function,
+                samplewise_center=True,
+                validation_split=0.2)
+            evaluation_generator = data_generator.flow_from_dataframe(
+                self.training_labels,
+                directory=self.image_dir,
+                x_col="filename",
+                y_col=DATA_LABELS,
+                target_size=(self.image_width, self.image_height),
+                batch_size=1,
+                subset='validation',
+                shuffle=False,
+                seed=self.seed,
+                class_mode="other")
+        else:
+            print("Using evaluation set...")
+
+        prediction_manager = PredictionManager()
+
+        try:
+            self.model.predict_generator(
+                evaluation_generator,
+                callbacks=[prediction_manager],
+                verbose=1,
+                steps=len(evaluation_generator))
+        except KeyboardInterrupt:
+            print('\r\nJob interrupted (Ctrl+C): Outputting completed results')
+
+        predictions = prediction_manager.getPredictions()
+
+        # Truncates evaluation generator if test was interrupted
+        evaluation_generator._data = evaluation_generator.data[:predictions.shape[0]]
+
+        self._save_prediction_histograms(predictions)
+        predictions = predictions > threshold
+        ground_truth = evaluation_generator.data
+        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(evaluation_generator, predictions, best=True)
+        stats["top_10_worst"] = self._top_images(evaluation_generator, predictions, best=False)
+        return stats
+
+    def _top_images(self, labels, predictions, top_k=10, best=False):
+        """Collect top-k best or top-k worst predicted images
+
+        Keyword arguments:
+        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(labels.data).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(labels.filenames[i])
+
+        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")

requirements.txt

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

setup.py

@@ -1,21 +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=[],
-)

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

utils/reroll.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+
+import csv
+import random
+
+with open("newlabels.csv", "w", newline="") as csvfile:
+    # Open the old labels file
+    with open("labels.csv", newline="") as oldfile:
+        # Get the header labels
+        csvreader = csv.DictReader(oldfile)
+        fieldnames = next(csvreader)
+        labelwriter = csv.DictWriter(csvfile, fieldnames=fieldnames)
+        labelwriter.writeheader()
+
+        # Loop through the file
+        rows = []
+        for row in csvreader:
+            row["evaluation"] = random.random() > 0.8
+            rows.append(row)
+        labelwriter.writerows(rows)
+print("Made new labels file: newlabels.csv")

utils/resizer.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+# Converts full-size images to a given size (replacing the originals)
+for f in images_full_size/*
+do
+	convert -resize 224x224 $f $f
+done