Saving

Works now!

eyeballer/__init__.py

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

eyeballer/augmentation.py

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

eyeballer/eyeballer.py

@@ -1,6 +1,9 @@
 #!/usr/bin/env python3
 import click
-from model import EyeballModel
+import csv
+
+from model import EyeballModel, DATA_LABELS
+from visualization import HeatMap
 
 
 @click.group(invoke_without_command=True)
@@ -31,24 +34,39 @@ def train(ctx, graphs, batchsize, epochs):
 
 
 @cli.command()
-@click.argument('screenshot', nargs=-1)  # Nargs=-1 means you can pass as many in as you want
+@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'])
-    for fn in screenshot:
-        model.predict(fn)
+    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:
+        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")
 
 
 def pretty_print_evaluation(results):
-    """Print a human-readable summary of the evaulation results"""
-    print("Custom404 Precision Score: " + str(round(results["custom404"]["precision"] * 100, 2)) + "%")
-    print("Custom404 Recall Score: " + str(round(results["custom404"]["recall"] * 100, 2)) + "%")
-    print("Login Precision Score: " + str(round(results["login"]["precision"] * 100, 2)) + "%")
-    print("Login Recall Score: " + str(round(results["login"]["recall"] * 100, 2)) + "%")
-    print("Homepage Precision Score: " + str(round(results["homepage"]["precision"] * 100, 2)) + "%")
-    print("Homepage Recall Score: " + str(round(results["homepage"]["recall"] * 100, 2)) + "%")
-    print("Overall Binary Accuracy: " + str(round(results["total_binary_accuracy"] * 100, 2)) + "%")
-    print("All or nothing accuracy: " + str(round(results["all_or_nothing_accuracy"] * 100, 2)) + "%")
+    """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]}")
 
 
 @cli.command()

eyeballer/model.py

@@ -1,10 +1,10 @@
 import os
 import random
-import time
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import Augmentor
 
 from keras import regularizers
 from keras.applications import MobileNet
@@ -16,9 +16,10 @@ from keras.models import Model
 from keras.optimizers import Adam
 from keras.preprocessing import image
 from keras.preprocessing.image import ImageDataGenerator
-from sklearn.metrics import classification_report, accuracy_score, jaccard_similarity_score
+from sklearn.metrics import classification_report, accuracy_score, hamming_loss
+from augmentation import EyeballerAugmentation
 
-DATA_LABELS = ["custom404", "login", "homepage"]
+DATA_LABELS = ["custom404", "login", "homepage", "oldlooking"]
 
 
 class EyeballModel:
@@ -26,7 +27,8 @@ class EyeballModel:
 
     Contains high-level functions for training, evaluating, and predicting.
     """
-    checkpoint_file = "weights-snapshot.h5"
+    graphs_directory = "graphs/"
+    checkpoint_file = "weights.h5"
     image_dir = "images_224x224/"
     image_width, image_height = 224, 224
     input_shape = (image_width, image_height, 3)
@@ -47,11 +49,11 @@ class EyeballModel:
         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)
+        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
+        #     layer.trainable = False
 
         adam = Adam(lr=0.0005)
         self.model.compile(optimizer=adam, loss="binary_crossentropy", metrics=["accuracy"])
@@ -81,6 +83,20 @@ class EyeballModel:
         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.
 
@@ -91,15 +107,10 @@ class EyeballModel:
         """
         print("Training with seed: " + str(self.seed))
 
-        # Data augmentation
         data_generator = ImageDataGenerator(
-            preprocessing_function=preprocess_input,
-            shear_range=0.0,
-            zoom_range=0.2,
-            brightness_range=(0.7, 1.3),
-            samplewise_center=True,
+            preprocessing_function=self.preprocess_training_function,
             validation_split=0.2,
-            horizontal_flip=False)
+            samplewise_center=True)
 
         training_generator = data_generator.flow_from_dataframe(
             self.training_labels,
@@ -138,6 +149,8 @@ class EyeballModel:
             verbose=1)
 
         if print_graphs:
+            if not os.path.exists(self.graphs_directory):
+                os.makedirs(self.graphs_directory)
             # Plot training & validation accuracy values
             plt.plot(history.history['acc'])
             plt.plot(history.history['val_acc'])
@@ -145,7 +158,7 @@ class EyeballModel:
             plt.ylabel('Accuracy')
             plt.xlabel('Epoch')
             plt.legend(['Train', 'Validation'], loc='upper left')
-            plt.savefig("accuracy.png")
+            plt.savefig(self.graphs_directory + "accuracy.png")
             plt.clf()
             plt.cla()
             plt.close()
@@ -157,28 +170,60 @@ class EyeballModel:
             plt.ylabel('Loss')
             plt.xlabel('Epoch')
             plt.legend(['Train', 'Validation'], loc='upper left')
-            plt.savefig("loss.png")
+            plt.savefig(self.graphs_directory + "loss.png")
 
-    def predict(self, filename):
+    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
 
         Keyword arguments:
-        filename -- The path to the file that we'll be evaluating
+        path -- 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)
+        # 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
 
-        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)
+        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
@@ -203,13 +248,9 @@ class EyeballModel:
             print("Using validation set...")
             # Data augmentation
             data_generator = ImageDataGenerator(
-                preprocessing_function=preprocess_input,
-                shear_range=0.0,
-                zoom_range=0.2,
-                brightness_range=(0.7, 1.3),
+                preprocessing_function=self.preprocess_training_function,
                 samplewise_center=True,
-                validation_split=0.2,
-                horizontal_flip=False)
+                validation_split=0.2)
             evaluation_generator = data_generator.flow_from_dataframe(
                 self.training_labels,
                 directory=self.image_dir,
@@ -225,13 +266,61 @@ class EyeballModel:
             print("Using evaluation set...")
 
         predictions = self.model.predict_generator(evaluation_generator, verbose=1, steps=len(evaluation_generator))
+        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"] = jaccard_similarity_score(ground_truth, predictions)
+        stats["total_binary_accuracy"] = 1 - hamming_loss(ground_truth, predictions)
         stats["all_or_nothing_accuracy"] = accuracy_score(ground_truth, predictions)
+        stats["top_10_best"] = self._top_images(evaluation_generator, predictions, best=True)
+        stats["top_10_worst"] = self._top_images(evaluation_generator, predictions, best=False)
         return stats
 
-    def get_data_column(self, data_slice, data):
-        return np.reshape(np.array(data)[data_slice], len(data)).tolist()
+    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")

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

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