Prediction: Made top results configurable

2 files changed, 18 insertions, 4 deletions

diff --git a/file-tagger.py b/file-tagger.py
index 2e4fe40..a36b602 100644
--- a/file-tagger.py
+++ b/file-tagger.py
@@ -112,7 +112,7 @@ def walk(args):
                     array = np.expand_dims(array_pre, axis=0)
                     array = preprocess_input(array)
                     predictions = model.predict(array)
-                    classes = decode_predictions(predictions, top=10)
+                    classes = decode_predictions(predictions, top=args["predict_images_top"])
                     logger.debug("Predicted image classes: {}".format(classes[0]))
                     tags.update([name for _, name, _ in classes[0]])
                     array_pre = cv2.rotate(array_pre, cv2.ROTATE_90_CLOCKWISE)
@@ -153,6 +153,7 @@ if __name__ == "__main__":
     parser.add_argument('-b', '--base', nargs='?', default='./test', type=dir_path, help='Base directory for walking (default: %(default)s)')
     parser.add_argument('-g', '--gui', nargs='?', const=1, default=False, type=bool, help='Show main GUI (default: %(default)s)')
     parser.add_argument('--predict-images', nargs='?', const=1, default=False, type=bool, help='Use prediction for image tagging (default: %(default)s)')
+    parser.add_argument('--predict-images-top', nargs='?', const=1, default=10, type=int, help='Defines how many top prediction keywords should be used (default: %(default)s)')
     parser.add_argument('--gui-tag', nargs='?', const=1, default=False, type=bool, help='Show GUI for tagging (default: %(default)s)')
     parser.add_argument('--open-system', nargs='?', const=1, default=False, type=bool, help='Open all files with system default (default: %(default)s)')
     parser.add_argument('-v', '--verbose', action="count", default=0, help="Verbosity level")
@@ -172,6 +173,7 @@ if __name__ == "__main__":
         "base": args.base,
         "gui": args.gui,
         "predict_images": args.predict_images,
+        "predict_images_top": args.predict_images_top,
         "gui_tag": args.gui_tag,
         "open_system": args.open_system,
         "verbosity": args.verbose
diff --git a/gui.py b/gui.py
index 947b3d4..248bbe0 100644
--- a/gui.py
+++ b/gui.py
@@ -11,22 +11,28 @@ class GuiMain(object):
         self.__args = args
         self.__base = StringVar(self.__master, value=args["base"])
         self.__predict_images = BooleanVar(self.__master, value=args["predict_images"])
+        self.__predict_images_top = StringVar(self.__master, value=str(args["predict_images_top"]))
         self.__gui_tag = BooleanVar(self.__master, value=args["gui_tag"])
         self.__open_system = BooleanVar(self.__master, value=args["open_system"])
 
+        validate_number = (self.__master.register(self.__validate_number))
+
         Label(self.__master, text="Base directory for walking:").grid(row=0, column=0)
         Entry(self.__master, textvariable=self.__base).grid(row=0, column=1, columnspan=2)
         Button(self.__master, text="Browse", command=self.__browse).grid(row=0, column=3)
         Checkbutton(self.__master, text="Use prediction for image tagging", variable=self.__predict_images).grid(row=1, column=0, columnspan=4, sticky=W)
-        Checkbutton(self.__master, text="Show GUI for tagging", variable=self.__gui_tag).grid(row=2, column=0, columnspan=4, sticky=W)
-        Checkbutton(self.__master, text="Open all files with system default", variable=self.__open_system).grid(row=3, column=0, columnspan=4, sticky=W)
-        Button(self.__master, text="Start", command=self.__master.destroy).grid(row=4, column=0, columnspan=4)
+        Label(self.__master, text="Number of top results:").grid(row=2, column=0)
+        Entry(self.__master, textvariable=self.__predict_images_top, validate='all', validatecommand=(validate_number, '%P')).grid(row=2, column=1, columnspan=1)
+        Checkbutton(self.__master, text="Show GUI for tagging", variable=self.__gui_tag).grid(row=3, column=0, columnspan=4, sticky=W)
+        Checkbutton(self.__master, text="Open all files with system default", variable=self.__open_system).grid(row=4, column=0, columnspan=4, sticky=W)
+        Button(self.__master, text="Start", command=self.__master.destroy).grid(row=5, column=0, columnspan=4)
 
     def loop(self):
         self.__master.mainloop()
 
         self.__args["base"] = self.__base.get()
         self.__args["predict_images"] = self.__predict_images.get()
+        self.__args["predict_images_top"] = int(self.__predict_images_top.get())
         self.__args["gui_tag"] = self.__gui_tag.get()
         self.__args["open_system"] = self.__open_system.get()
         return self.__args
@@ -35,6 +41,12 @@ class GuiMain(object):
         filename = filedialog.askdirectory()
         self.__base.set(filename)
 
+    def __validate_number(self, P):
+        if str.isdigit(P) or P == "":
+            return True
+        else:
+            return False
+
 class GuiTag(object):
     RETURN_NEXT = 0
     RETURN_ABORT = 1