ml_quick_label.py

   1 #!/usr/bin/env python3
   2
   3 import glob
   4 import os
   5 from typing import Callable, List, NamedTuple, Set
   6
   7 import argparse_utils
   8 import config
   9 import input_utils
  10
  11 parser = config.add_commandline_args(
  12     f"ML Quick Labeler ({__file__})",
  13     "Args related to quick labeling of ML training data",
  14 )
  15 parser.add_argument(
  16     "--ml_quick_label_skip_list_path",
  17     default="./qlabel_skip_list.txt",
  18     metavar="FILENAME",
  19     type=argparse_utils.valid_filename,
  20     help="Path to file in which to store already labeled data",
  21 )
  22
  23
  24 class InputSpec(NamedTuple):
  25     image_file_glob: str
  26     image_file_to_features_file: Callable[[str], str]
  27     label: str
  28     valid_keystrokes: List[str]
  29     prompt: str
  30     keystroke_to_label: Callable[[str], str]
  31
  32
  33 def read_skip_list() -> Set[str]:
  34     ret: Set[str] = set()
  35     quick_skip_file = config.config['ml_quick_label_skip_list_path']
  36     if not os.path.exists(quick_skip_file):
  37         return ret
  38     with open(quick_skip_file, 'r') as f:
  39         lines = f.readlines()
  40     for line in lines:
  41         line = line[:-1]
  42         line.strip()
  43         ret.add(line)
  44     return ret
  45
  46
  47 def write_skip_list(skip_list) -> None:
  48     quick_skip_file = config.config['ml_quick_label_skip_list_path']
  49     with open(quick_skip_file, 'w') as f:
  50         for filename in skip_list:
  51             filename = filename.strip()
  52             if len(filename) > 0:
  53                 f.write(f'{filename}\n')
  54
  55
  56 def label(in_spec: InputSpec) -> None:
  57     images = glob.glob(in_spec.image_file_glob)
  58
  59     skip_list = read_skip_list()
  60     for image in images:
  61         if image in skip_list:
  62             continue
  63         features = in_spec.image_file_to_features_file(image)
  64         if features is None or not os.path.exists(features):
  65             continue
  66
  67         # Render features and image.
  68         with open(features, "r") as f:
  69             lines = f.readlines()
  70         skip = False
  71         for line in lines:
  72             line = line[:-1]
  73             if in_spec.label in line:
  74                 skip = True
  75         if skip:
  76             skip_list.add(image)
  77             continue
  78
  79         os.system(f'xv {image} &')
  80         keystroke = input_utils.single_keystroke_response(
  81             in_spec.valid_keystrokes,
  82             prompt=in_spec.prompt,
  83         )
  84         os.system('killall xv')
  85
  86         label_value = in_spec.keystroke_to_label(keystroke)
  87         with open(features, "a") as f:
  88             f.write(f"{in_spec.label}: {label_value}\n")
  89         skip_list.add(image)
  90
  91     write_skip_list(skip_list)