Stats are now generated in a saner way,

Label names come from a global configurable (likely move this to a
config file later)
Stats are outputted as a dict.

eyeballer/eyeballer.py

@@ -22,6 +22,7 @@ from keras.preprocessing.image import ImageDataGenerator
 from sklearn.metrics import precision_score, recall_score
 
 # Parse the arguments
+DATA_LABELS = ["custom404", "login", "homepage"]
 parser = argparse.ArgumentParser(description='Give those screenshots of yours a quick eyeballing')
 parser.add_argument("mode", help="Mode of operation. One of \"train\", \"predict\", or \"evaluate\".")
 parser.add_argument("--weights", help="Weights file for input/output")
@@ -36,12 +37,23 @@ args = parser.parse_args()
 
 
 class EyeballModel:
+    """The primary model class of Eyeballer.
+
+    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)
 
     def __init__(self, print_summary=False, weights_file="weights.h5", seed=0):
+        """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
@@ -79,8 +91,14 @@ class EyeballModel:
         else:
             print("No model loaded from file")
 
-    # Training
     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 augmentation
@@ -88,6 +106,7 @@ class EyeballModel:
             preprocessing_function=preprocess_input,
             shear_range=0.0,
             zoom_range=0.2,
+            brightness_range=(0.7, 1.3),
             samplewise_center=True,
             validation_split=0.2,
             horizontal_flip=False)
@@ -96,7 +115,7 @@ class EyeballModel:
             self.training_labels,
             directory=self.image_dir,
             x_col="filename",
-            y_col=["custom404", "login", "homepage"],
+            y_col=DATA_LABELS,
             target_size=(self.image_width, self.image_height),
             batch_size=batch_size,
             subset='training',
@@ -107,7 +126,7 @@ class EyeballModel:
             self.training_labels,
             directory=self.image_dir,
             x_col="filename",
-            y_col=["custom404", "login", "homepage"],
+            y_col=DATA_LABELS,
             target_size=(self.image_width, self.image_height),
             batch_size=batch_size,
             subset='validation',
@@ -150,8 +169,12 @@ class EyeballModel:
             plt.legend(['Train', 'Validation'], loc='upper left')
             plt.savefig("loss.png")
 
-    # Single predict on a file
     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)
@@ -167,8 +190,12 @@ class EyeballModel:
         print("\tHome Page: " + str(round(prediction[0][2] * 100, 2)) + "%")
         print("Prediction Took (seconds): ", time_end - time_start)
 
-    # Evaluate performance against the persistent evaluation data set
     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)
@@ -176,7 +203,7 @@ class EyeballModel:
             self.evaluation_labels,
             directory=self.image_dir,
             x_col="filename",
-            y_col=["custom404", "login", "homepage"],
+            y_col=DATA_LABELS,
             target_size=(self.image_width, self.image_height),
             shuffle=False,
             batch_size=1,
@@ -189,6 +216,7 @@ class EyeballModel:
                 preprocessing_function=preprocess_input,
                 shear_range=0.0,
                 zoom_range=0.2,
+                brightness_range=(0.7, 1.3),
                 samplewise_center=True,
                 validation_split=0.2,
                 horizontal_flip=False)
@@ -196,7 +224,7 @@ class EyeballModel:
                 self.training_labels,
                 directory=self.image_dir,
                 x_col="filename",
-                y_col=["custom404", "login", "homepage"],
+                y_col=DATA_LABELS,
                 target_size=(self.image_width, self.image_height),
                 batch_size=1,
                 subset='validation',
@@ -228,34 +256,18 @@ class EyeballModel:
                     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()
+        stats = {label: dict() for label in DATA_LABELS}
+        for index, scores in enumerate(stats.values()):
+            data_slice = np.index_exp[:, index:index+1]  # variable is verbosely named to avoid shadowing np.slice
+            labels_column = self.get_data_column(data_slice, labels)
+            predictions_column = self.get_data_column(data_slice, predictions)
+            scores['Precision'] = precision_score(labels_column, predictions_column)
+            scores['Recall'] = recall_score(labels_column, predictions_column)
 
-        login_labels = np.reshape(np.array(labels)[:, 1:2], len(labels)).tolist()
-        login_predicitons = np.reshape(np.array(predictions)[:, 1:2], len(predictions)).tolist()
+        return stats
 
-        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)) + "%")
+    def get_data_column(self, data_slice, data):
+        return np.reshape(np.array(data)[data_slice], len(data)).tolist()
 
 
 seed = args.seed

setup.py

@@ -18,4 +18,5 @@ setup(
         'pillow'
     ],
     test_requires=[],
+    python_requires='>=3.6',    # We use some python 3.6+ only features, like ordered dicts by default and f-strings
 )

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