Created Cli using click.

If seed is unspecified on the cli, it will be randomly generated.

eyeballer/__init__.py

@@ -0,0 +1 @@
+from eyeballer.eyeballer import EyeballModel

eyeballer/eyeballer.py

@@ -21,27 +21,25 @@ from keras.preprocessing import image
 from keras.preprocessing.image import ImageDataGenerator
 from sklearn.metrics import precision_score, recall_score
 
-# Parse the arguments
-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")
-parser.add_argument("--batchsize", help="Batch size", default=32, type=int)
-parser.add_argument("--epochs", help="Number of epochs", default=20, type=int)
-parser.add_argument("--screenshot", help="File to predict")
-parser.add_argument("--summary", help="Print model summary at start", action='store_true')
-parser.add_argument("--graphs", help="Save accuracy and loss graphs to file", action='store_true')
-parser.add_argument("--seed", help="RNG seed for data shuffling and transformations", default=0, type=int)
-
-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):
+    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
@@ -69,7 +67,10 @@ class EyeballModel:
 
         # Shuffle the training labels
         self.seed = seed
-        print("Using seed: " + str(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)
 
@@ -79,8 +80,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
@@ -150,8 +157,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 +178,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)
@@ -182,7 +197,7 @@ class EyeballModel:
             batch_size=1,
             class_mode="other")
         # If a seed was selected, then also evaluate on the validation set for that seed
-        if args.seed > 0:
+        if self.seed is not None:
             print("Using validation set...")
             # Data augmentation
             data_generator = ImageDataGenerator(
@@ -257,20 +272,3 @@ class EyeballModel:
         print("All or nothing accuracy: " + str(
             round((all_or_nothing_success / len(evaluation_generator)) * 100, 2)) + "%")
 
-
-seed = args.seed
-if args.seed == 0:
-    seed = random.randint(0, 999999)
-
-model = EyeballModel(weights_file=args.weights, print_summary=args.summary, seed=seed)
-if args.mode == "train":
-    model.train(print_graphs=args.graphs, batch_size=args.batchsize, epochs=args.epochs)
-elif args.mode == "predict":
-    if args.screenshot is None:
-        print("Error: Use the --screenshot flag to specify what file you want to predict on")
-    else:
-        model.predict(args.screenshot)
-elif args.mode == "evaluate":
-    model.evaluate(threshold=0.5)
-else:
-    print("Error: " + args.mode + " is not a valid mode. Try one of \"train\", \"predict\", or \"evaluate\"")

eyeballer/main.py

@@ -0,0 +1,58 @@
+#!/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()
+

setup.py

@@ -15,7 +15,7 @@ setup(
         'numpy',
         'pandas',
         'matplotlib',
-        'pillow'
+        'pillow', 'click'
     ],
     test_requires=[],
 )