summaryrefslogtreecommitdiff
path: root/ml/quick_label.py
blob: da9a1d2e00d15e32cce697e3782a5cfce87ea584 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
#!/usr/bin/env python3

"""A helper to facilitate quick manual labeling of ML training data."""

import glob
import logging
import os
import warnings
from dataclasses import dataclass
from typing import Callable, List, Optional, Set

import argparse_utils
import config

logger = logging.getLogger(__name__)
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.",
)
parser.add_argument(
    "--ml_quick_label_use_skip_lists",
    default=True,
    action=argparse_utils.ActionNoYes,
    help='Should we use a skip list file to speed up execution?',
)
parser.add_argument(
    "--ml_quick_label_overwrite_labels",
    default=False,
    action=argparse_utils.ActionNoYes,
    help='Enable overwriting existing labels; default is to not relabel.',
)


@dataclass
class InputSpec:
    """A wrapper around the input data we need to operate; should be
    populated by the caller."""

    image_file_glob: Optional[str] = None
    image_file_prepopulated_list: Optional[List[str]] = None
    image_file_to_features_file: Optional[Callable[[str], str]] = None
    label: str = ''
    valid_keystrokes: List[str] = []
    prompt: str = ''
    keystroke_to_label: Optional[Callable[[str], str]] = None


def read_skip_list() -> Set[str]:
    ret: Set[str] = set()
    if config.config['ml_quick_label_use_skip_lists']:
        quick_skip_file = config.config['ml_quick_label_skip_list_path']
        if os.path.exists(quick_skip_file):
            with open(quick_skip_file, 'r') as f:
                lines = f.readlines()
            for line in lines:
                line = line[:-1]
                line.strip()
                ret.add(line)
        logger.debug('Read %s and found %d entries.', quick_skip_file, len(ret))
    return ret


def write_skip_list(skip_list) -> None:
    if config.config['ml_quick_label_use_skip_lists']:
        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')
        logger.debug('Updated %s', quick_skip_file)


def label(in_spec: InputSpec) -> None:
    import input_utils

    images = []
    if in_spec.image_file_glob is not None:
        images += glob.glob(in_spec.image_file_glob)
    elif in_spec.image_file_prepopulated_list is not None:
        images += in_spec.image_file_prepopulated_list
    else:
        raise ValueError('One of image_file_glob or image_file_prepopulated_list is required')

    skip_list = read_skip_list()
    for image in images:
        if image in skip_list:
            logger.debug('Skipping %s because of the skip list', image)
            continue
        assert in_spec.image_file_to_features_file
        features = in_spec.image_file_to_features_file(image)
        if features is None or not os.path.exists(features):
            msg = f'File {image} yielded file {features} which does not exist, SKIPPING.'
            logger.warning(msg)
            warnings.warn(msg)
            continue

        # Render features and image.
        filtered_lines = []
        with open(features, "r") as f:
            lines = f.readlines()
        saw_label = False
        for line in lines:
            line = line[:-1]
            if in_spec.label not in line:
                filtered_lines.append(line)
            else:
                saw_label = True

        if not saw_label or config.config['ml_quick_label_overwrite_labels']:
            logger.info(features)
            os.system(f'xv {image} &')
            keystroke = input_utils.single_keystroke_response(
                in_spec.valid_keystrokes,
                prompt=in_spec.prompt,
            )
            os.system('killall xv')
            assert in_spec.keystroke_to_label
            label_value = in_spec.keystroke_to_label(keystroke)
            filtered_lines.append(f"{in_spec.label}: {label_value}\n")
            with open(features, 'w') as f:
                f.writelines(line + '\n' for line in filtered_lines)
            skip_list.add(image)
    write_skip_list(skip_list)