Fixed progressbar.progressbar issue

- All prediction tests

.gitignore

@@ -116,3 +116,5 @@ venv.bak/
 results.html
 !requirements.txt
 !requirements-gpu.txt
+!tests/models/*
+!tests/data/*

README.md

@@ -19,19 +19,19 @@ sudo pip3 install -r requirements-gpu.txt
 To eyeball some screenshots, just run the "predict" mode:
 
 ```
-eyeballer --weights YOUR_WEIGHTS.h5 predict YOUR_FILE.png
+eyeballer.py --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/
+eyeballer.py --weights YOUR_WEIGHTS.h5 predict PATH_TO/YOUR_FILES/
 ```
 
 ## Training
 To train a new model, run:
 ```
-eyeballer train
+eyeballer.py 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.
@@ -43,5 +43,5 @@ This will output a new model file (weights.h5 by default).
 You just trained a new model, cool! Let's see how well it performs against some images it's never seen before, across a variety of metrics:
 
 ```
-eyeballer --weights YOUR_WEIGHTS.h5 evaluate
+eyeballer.py --weights YOUR_WEIGHTS.h5 evaluate
 ```

eyeballer.py

@@ -46,7 +46,7 @@ def predict(ctx, screenshot, heatmap, threshold):
 
     if heatmap:
         # Generate a heatmap
-        HeatMap(screenshot, model).generate()
+        HeatMap(screenshot, model, threshold).generate()
 
     if not results:
         print("Error: Input file does not exist")
@@ -60,11 +60,11 @@ def predict(ctx, screenshot, heatmap, threshold):
             labelwriter.writerows(results)
 
         print("Output written to results.csv")
-        buildHTML(processResults(results))
+        buildHTML(processResults(results, threshold))
         print("HTML written to results.html")
 
 
-def processResults(results):
+def processResults(results, threshold):
     '''Filter the initial results dictionary and reformat it for use in JS.
 
         Keyword arguments:
@@ -79,8 +79,7 @@ def processResults(results):
         for label, label_info in result.items():
             if (label == 'filename'):
                 pass
-
-            elif label_info:
+            elif label_info > threshold:
                 positiveTags.append(label)
 
         jsResults[result['filename']] = positiveTags

eyeballer/model.py

@@ -13,7 +13,7 @@ import pandas as pd
 import Augmentor
 import tensorflow as tf
 
-from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
+from tensorflow.keras.applications.mobilenet import preprocess_input
 from sklearn.metrics import classification_report, accuracy_score, hamming_loss
 from eyeballer.augmentation import EyeballerAugmentation
 
@@ -27,7 +27,7 @@ class EyeballModel:
     """
     graphs_directory = "graphs/"
     checkpoint_file = "weights.h5"
-    image_dir = "images_224x224/"
+    image_dir = "images/"
     image_width, image_height = 224, 224
     input_shape = (image_width, image_height, 3)
 
@@ -41,7 +41,7 @@ class EyeballModel:
         """
         # # Build the model
         self.model = tf.keras.Sequential()
-        pretrained_layer = tf.keras.applications.mobilenet_v2.MobileNetV2(weights='imagenet', include_top=False, input_shape=self.input_shape)
+        pretrained_layer = tf.keras.applications.mobilenet.MobileNet(weights='imagenet', include_top=False, input_shape=self.input_shape)
         self.model.add(pretrained_layer)
         self.model.add(tf.keras.layers.GlobalAveragePooling2D())
         self.model.add(tf.keras.layers.Dense(256, activation="relu"))
@@ -141,8 +141,12 @@ class EyeballModel:
             class_mode="other")
 
         # Model checkpoint - Saves model weights when validation accuracy improves
-        callbacks = [tf.keras.callbacks.ModelCheckpoint(self.checkpoint_file, monitor='val_loss', verbose=1, save_best_only=True,
-                     save_weights_only=True, mode='min')]
+        callbacks = [tf.keras.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,
@@ -212,7 +216,8 @@ class EyeballModel:
             raise FileNotFoundError
 
         results = []
-        for screenshot in progressbar.progressbar(screenshots):
+        pbar = progressbar.ProgressBar()
+        for screenshot in pbar(screenshots):
             # Load the image into memory
             img = None
             try:
@@ -223,6 +228,9 @@ class EyeballModel:
             except IsADirectoryError:
                 print("\nWARN: Skipping directory: ", screenshot)
                 continue
+            except OSError:
+                print("\nWARN: Skipping empty or corrupt file: ", screenshot)
+                continue
 
             prediction = self.model.predict(img, batch_size=1)
             result = dict()

eyeballer/visualization.py

@@ -7,7 +7,7 @@ from eyeballer.model import DATA_LABELS
 
 
 class HeatMap():
-    def __init__(self, screenshot_path, model, boxsize=28, step=7):
+    def __init__(self, screenshot_path, model, threshold=0.5, boxsize=28, step=7):
         self.model = model
 
         img = image.load_img(screenshot_path, target_size=(self.model.image_width, self.model.image_height))
@@ -20,6 +20,7 @@ class HeatMap():
         self.y = 0
         self.screenshot_path = screenshot_path
         self.gamma = 3
+        self.threshold = threshold
 
     def generate(self, output_file="heatmap.png"):
         """ Make a single heatmap image and return it """
@@ -30,11 +31,11 @@ class HeatMap():
         for label in DATA_LABELS:
             boxsize = self.boxsize
             # Ignore this label if it didn't predict positively (true label)
-            if results[label] > 0.5:
+            if results[label] > self.threshold:
                 worst_score = 1
-                while worst_score > 0.5:
+                while worst_score > self.threshold:
                     heatmap, worst_score = self._get_heatmap(label, boxsize)
-                    if worst_score > 0.5:
+                    if worst_score > self.threshold:
                         boxsize += 28
                         print("Didn't get a good image for {}. Trying again with a bigger boxsize: {}".format(label, boxsize))
                     else:

tests/PredictTest.py

@@ -1,15 +1,22 @@
 import errno
 import os
 import unittest
+import sys
 
 from eyeballer.model import EyeballModel
 
 
 class PredictTest(unittest.TestCase):
     def setUp(self):
-        weights_file = "latest_weights.h5"
+
+        class DummyFile(object):
+            def write(self, x): pass
+            def flush(self): pass
+
+        sys.stdout = DummyFile()
+        weights_file = "tests/models/test_weights.h5"
         if not os.path.isfile(weights_file):
-            print("Error: Symlink the latest weights file to 'latest_weights.h5'")
+            print("Error: Symlink the latest weights file to " + weights_file)
             raise FileNotFoundError(
                 errno.ENOENT,
                 os.strerror(errno.ENOENT),
@@ -30,7 +37,7 @@ class PredictTest(unittest.TestCase):
         }
         same_seed_model = EyeballModel(**model_kwargs)
 
-        screenshot = "tests/data/test.png"
+        screenshot = "tests/data/404.png"
 
         results_one = self.model.predict(screenshot)
         results_two = same_seed_model.predict(screenshot)
@@ -45,40 +52,77 @@ class PredictTest(unittest.TestCase):
         }
         same_seed_model = EyeballModel(**model_kwargs)
 
-        screenshot = "tests/data/test.png"
+        screenshot = "tests/data/404.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
+        self.assertEqual(results_one, results_two)
 
     def test_predict_custom404(self):
-        screenshot = "tests/data/custom404.png"
-        results = self.model.predict(screenshot)
-        self.assertEqual(self.get_prediction_by_name(results), "custom404")
+        screenshot = "tests/data/404.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["custom404"], 0.5)
+
+    def test_predict_not_custom404(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertLess(results["custom404"], 0.5)
 
     def test_predict_login(self):
         screenshot = "tests/data/login.png"
-        results = self.model.predict(screenshot)
-        self.assertEqual(self.get_prediction_by_name(results), "login")
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["login"], 0.5)
+
+    def test_predict_not_login(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        print(screenshot, results)
+        self.assertLess(results["login"], 0.5)
 
     def test_predict_homepage(self):
         screenshot = "tests/data/homepage.png"
-        results = self.model.predict(screenshot)
-        self.assertEqual(self.get_prediction_by_name(results), "homepage")
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["homepage"], 0.5)
+
+    def test_predict_not_homepage(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertLess(results["homepage"], 0.5)
 
     def test_predict_oldlooking(self):
-        screenshot = "tests/data/oldlooking.png"
-        results = self.model.predict(screenshot)
-        self.assertEqual(self.get_prediction_by_name(results), "oldlooking")
+        screenshot = "tests/data/old-looking.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertGreater(results["oldlooking"], 0.5)
+
+    def test_predict_not_oldlooking(self):
+        screenshot = "tests/data/nothing.png"
+        results = self.model.predict(screenshot)[0]
+        self.assertLess(results["oldlooking"], 0.5)
+
+    def test_file_doesnt_exist(self):
+        screenshot = "tests/data/doesnotexist.png"
+        try:
+            self.model.predict(screenshot)[0]
+            self.fail("FileNotFoundError was expected but not found")
+        except FileNotFoundError:
+            pass
+
+    def test_folder(self):
+        screenshots = "tests/data/"
+        results = self.model.predict(screenshots)
+        self.assertEqual(len(results), 7)
+
+    def test_file_is_empty(self):
+        """
+        We're just testing that it doesn't crash, basically
+        """
+        screenshot = "tests/data/empty.png"
+        self.model.predict(screenshot)
+
+    def test_file_is_invalid(self):
+        """
+        We're just testing that it doesn't crash, basically
+        """
+        screenshot = "tests/data/invalid.png"
+        self.model.predict(screenshot)

tests/data/invalid.png

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