diff --git a/bus_analyzer/bus_analyzer/extract.py b/bus_analyzer/bus_analyzer/extract.py index 87c968f..e221791 100644 --- a/bus_analyzer/bus_analyzer/extract.py +++ b/bus_analyzer/bus_analyzer/extract.py @@ -245,16 +245,22 @@ def read_frame(frames, prototypes_path="./prototypes", verbose=False): value, score, digits = recognize_odometer(prototypes, frame) if verbose: - for guess, score, all_scores in digits: - print("Digit = {} with score {}".format(guess, score)) + for guess, _score, all_scores in digits: + print("Digit = {} with score {}".format(guess, _score)) print("{}: odo {} with score {}".format(filename, value, score)) value, score, digits = recognize_clock(prototypes, frame) if verbose: - for guess, score, all_scores in digits: - print("Digit = {} with score {}".format(guess, score)) + for guess, _score, all_scores in digits: + print("Digit = {} with score {}".format(guess, _score)) print("{}: clock {} with score {}".format(filename, value, score)) + value, score, all_scores = recognize_time_of_day(frame) + if verbose: + for color, _score in all_scores: + print("{}: {}".format(color, _score)) + print("{}: time-of-day {} with score {}".format(filename, value, score)) + @cli def create_prototype(output, *images): @@ -281,9 +287,32 @@ def get_frame(*segments): print(filename) +def compare_colors(color_a, color_b): + MAX_ERROR_SQUARED = 255**2 * len(color_a) + return 1 - float(sum((a - b)**2 for a, b in zip(color_a, color_b))) / MAX_ERROR_SQUARED + + +def recognize_time_of_day(frame): + """Returns time of day, score, all scores.""" + COLORS = { + "day": (89, 236, 239), + "dusk": (199, 162, 205), + "night": (1, 1, 1), + "dawn": (51, 59, 142), + } + sample = frame.getpixel((177, 255)) + scores = [ + (tod, compare_colors(sample, color)) + for tod, color in COLORS.items() + ] + best, score = max(scores, key=lambda t: t[1]) + return best, score, scores + + def extract_segment(prototypes, segment): ODO_SCORE_THRESHOLD = 0.01 CLOCK_SCORE_THRESHOLD = 0.01 + TOD_SCORE_THRESHOLD = 0.9 frame_data = b"".join(extract_frame([segment], segment.start)) frame = Image.open(BytesIO(frame_data)) odometer, score, _ = recognize_odometer(prototypes, frame) @@ -292,7 +321,10 @@ def extract_segment(prototypes, segment): clock, score, _ = recognize_clock(prototypes, frame) if score < CLOCK_SCORE_THRESHOLD: clock = None - return odometer, clock + tod, score, _ = recognize_time_of_day(frame) + if score < TOD_SCORE_THRESHOLD: + tod = None + return odometer, clock, tod if __name__ == '__main__': diff --git a/bus_analyzer/bus_analyzer/main.py b/bus_analyzer/bus_analyzer/main.py index d5536df..1124bc2 100644 --- a/bus_analyzer/bus_analyzer/main.py +++ b/bus_analyzer/bus_analyzer/main.py @@ -25,8 +25,8 @@ def do_extract_segment(*segment_paths, prototypes_path="./prototypes"): prototypes = load_prototypes(prototypes_path) for segment_path in segment_paths: segment_info = parse_segment_path(segment_path) - odometer, clock = extract_segment(prototypes, segment_info) - print(f"{segment_path} {odometer} {clock}") + odometer, clock, tod = extract_segment(prototypes, segment_info) + print(f"{segment_path} {odometer} {clock} {tod}") @cli @@ -75,7 +75,7 @@ def compare_segments(dbconnect, base_dir='.', prototypes_path="./prototypes", si for old_odometer, segment in selected: path = os.path.join(base_dir, segment) segment_info = parse_segment_path(path) - odometer, clock = extract_segment(prototypes, segment_info) + odometer, clock, tod = extract_segment(prototypes, segment_info) results.append((segment, old_odometer, odometer)) matching = 0 @@ -112,24 +112,25 @@ def analyze_segment(conn, prototypes, segment_path, check_segment_name=None): assert segment_name == check_segment_name try: - odometer, clock = extract_segment(prototypes, segment_info) + odometer, clock, tod = extract_segment(prototypes, segment_info) except Exception: logging.warning(f"Failed to extract segment {segment_path!r}", exc_info=True) odometer = None error = traceback.format_exc() else: - logging.info(f"Got odometer = {odometer}, clock = {clock} for segment {segment_path!r}") + logging.info(f"Got odometer = {odometer}, clock = {clock}, time of day = {tod} for segment {segment_path!r}") error = None database.query( conn, """ - INSERT INTO bus_data (channel, timestamp, segment, error, odometer, clock) - VALUES (%(channel)s, %(timestamp)s, %(segment)s, %(error)s, %(odometer)s, %(clock)s) + INSERT INTO bus_data (channel, timestamp, segment, error, odometer, clock, timeofday) + VALUES (%(channel)s, %(timestamp)s, %(segment)s, %(error)s, %(odometer)s, %(clock)s, %(timeofday)s) ON CONFLICT (channel, timestamp, segment) DO UPDATE SET error = %(error)s, odometer = %(odometer)s, - clock = %(clock)s + clock = %(clock)s, + timeofday = %(timeofday)s """, channel=segment_info.channel, timestamp=segment_info.start, @@ -137,6 +138,7 @@ def analyze_segment(conn, prototypes, segment_path, check_segment_name=None): error=error, odometer=odometer, clock=clock, + timeofday=tod, ) diff --git a/postgres/setup.sh b/postgres/setup.sh index fca856d..eaa9fab 100644 --- a/postgres/setup.sh +++ b/postgres/setup.sh @@ -157,6 +157,7 @@ CREATE TABLE playlists ( -- be determined. -- The odometer column is in miles. The game shows the odometer to the 1/10th mile precision. -- The clock is in minutes since 00:00, in 12h time. +-- The time of day is one of "day", "dusk", "night", or "dawn" -- The segment may be NULL, which indicates a manually-inserted value. -- The primary key serves two purposes: -- It provides an index on channel, followed by a range index on timestamp @@ -170,6 +171,7 @@ CREATE TABLE bus_data ( error TEXT, odometer DOUBLE PRECISION, clock INTEGER, + timeofday TEXT, PRIMARY KEY (channel, timestamp, segment) ); EOSQL