1 files changed, 91 insertions, 0 deletions
diff --git a/ml_quick_label.py b/ml_quick_label.py
new file mode 100644
index 0000000..1c35982
--- /dev/null
+++ b/ml_quick_label.py
@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+
+import glob
+import os
+from typing import Callable, List, NamedTuple, Set
+
+import argparse_utils
+import config
+import input_utils
+
+parser = config.add_commandline_args(
+    f"ML Quick Labeler ({__file__})",
+    "Args related to quick labeling of ML training data",
+)
+parser.add_argument(
+    "--ml_quick_label_skip_list_path",
+    default="./qlabel_skip_list.txt",
+    metavar="FILENAME",
+    type=argparse_utils.valid_filename,
+    help="Path to file in which to store already labeled data",
+)
+
+
+class InputSpec(NamedTuple):
+    image_file_glob: str
+    image_file_to_features_file: Callable[[str], str]
+    label: str
+    valid_keystrokes: List[str]
+    prompt: str
+    keystroke_to_label: Callable[[str], str]
+
+
+def read_skip_list() -> Set[str]:
+    ret: Set[str] = set()
+    quick_skip_file = config.config['ml_quick_label_skip_list_path']
+    if not os.path.exists(quick_skip_file):
+        return ret
+    with open(quick_skip_file, 'r') as f:
+        lines = f.readlines()
+    for line in lines:
+        line = line[:-1]
+        line.strip()
+        ret.add(line)
+    return ret
+
+
+def write_skip_list(skip_list) -> None:
+    quick_skip_file = config.config['ml_quick_label_skip_list_path']
+    with open(quick_skip_file, 'w') as f:
+        for filename in skip_list:
+            filename = filename.strip()
+            if len(filename) > 0:
+                f.write(f'{filename}\n')
+
+
+def label(in_spec: InputSpec) -> None:
+    images = glob.glob(in_spec.image_file_glob)
+
+    skip_list = read_skip_list()
+    for image in images:
+        if image in skip_list:
+            continue
+        features = in_spec.image_file_to_features_file(image)
+        if features is None or not os.path.exists(features):
+            continue
+
+        # Render features and image.
+        with open(features, "r") as f:
+            lines = f.readlines()
+        skip = False
+        for line in lines:
+            line = line[:-1]
+            if in_spec.label in line:
+                skip = True
+        if skip:
+            skip_list.add(image)
+            continue
+
+        os.system(f'xv {image} &')
+        keystroke = input_utils.single_keystroke_response(
+            in_spec.valid_keystrokes,
+            prompt=in_spec.prompt,
+        )
+        os.system('killall xv')
+
+        label_value = in_spec.keystroke_to_label(keystroke)
+        with open(features, "a") as f:
+            f.write(f"{in_spec.label}: {label_value}\n")
+        skip_list.add(image)
+
+    write_skip_list(skip_list)