bus_analyzer: Move prototypes into subdir and allow non-digits

bus_analyzer/Dockerfile

@@ -13,7 +13,7 @@ RUN pip install /tmp/common && rm -r /tmp/common
 # Actual application
 RUN apk add ffmpeg
 COPY bus_analyzer /tmp/bus_analyzer
-RUN pip install /tmp/bus_analyzer && cp -r /tmp/bus_analyzer/odo-digit-prototypes / && rm -r /tmp/bus_analyzer
+RUN pip install /tmp/bus_analyzer && cp -r /tmp/bus_analyzer/prototypes / && rm -r /tmp/bus_analyzer
 
 LABEL org.opencontainers.image.source https://github.com/dbvideostriketeam/wubloader
 ENTRYPOINT ["python3", "-m", "bus_analyzer"]

bus_analyzer/bus_analyzer/extract.py

@@ -108,10 +108,11 @@ def recognize_digit(prototypes, image):
 	"""Takes a normalized digit image and returns (detected number, score, all_scores)
 	where score is between 0 and 1. Higher numbers are more certain the number is correct.
 	all_scores is for debugging.
+	If the most likely detection is NOT a number, None is returned instead.
 	"""
 	scores = sorted([
-		(compare_images(prototypes[n], image), n)
+		(compare_images(prototype, image), int(n) if n.isdigit() else None)
-		for n in range(10)
+		for n, prototype in prototypes["odo-digits"].items()
 	], reverse=True)
 	best_score, number = scores[0]
 	runner_up_score, _ = scores[1]
@@ -130,14 +131,20 @@ def compare_images(prototype, image):
 
 
 def load_prototypes(prototypes_path):
-	return [
+	prototypes = {}
-		Image.open(os.path.join(prototypes_path, "{}.png".format(n)))
+	for kind in os.listdir(prototypes_path):
-		for n in range(10)
+		prototypes[kind] = {}
-	]
+		path = os.path.join(prototypes_path, kind)
+		for filename in os.listdir(path):
+			if not filename.endswith(".png"):
+				continue
+			name = filename[:-4]
+			prototypes[kind][name] = Image.open(os.path.join(path, filename))
+	return prototypes
 
 
 @cli
-def read_digit(digit, prototypes_path="./odo-digit-prototypes", verbose=False):
+def read_digit(digit, prototypes_path="./prototypes", verbose=False):
 	"""For debugging. Compares an extracted digit image to each prototype and prints scores."""
 	prototypes = load_prototypes(prototypes_path)
 	digit = Image.open(digit)
@@ -165,7 +172,7 @@ def recognize_odometer(prototypes, frame):
 
 @cli
 @argh.arg("frames", nargs="+")
-def read_frame(frames, prototypes_path="./odo-digit-prototypes", verbose=False, include_last=False):
+def read_frame(frames, prototypes_path="./prototypes", verbose=False, include_last=False):
 	"""For testing. Takes any number of frame images (or segments) and prints the odometer reading."""
 	prototypes = load_prototypes(prototypes_path)
 	for filename in frames:

bus_analyzer/bus_analyzer/main.py

@@ -20,7 +20,7 @@ cli = argh.EntryPoint()
 
 @cli
 @argh.named("extract-segment")
-def do_extract_segment(*segment_paths, prototypes_path="./odo-digit-prototypes"):
+def do_extract_segment(*segment_paths, prototypes_path="./prototypes"):
 	"""Extract info from individual segments and print them"""
 	prototypes = load_prototypes(prototypes_path)
 	for segment_path in segment_paths:
@@ -30,7 +30,7 @@ def do_extract_segment(*segment_paths, prototypes_path="./odo-digit-prototypes")
 
 
 @cli
-def compare_segments(dbconnect, base_dir='.', prototypes_path="./odo-digit-prototypes", since=None, until=None, num=100, null_chance=0.25, verbose=False):
+def compare_segments(dbconnect, base_dir='.', prototypes_path="./prototypes", since=None, until=None, num=100, null_chance=0.25, verbose=False):
 	"""
 	Collect some representitive samples from the database and re-runs them to compare to previous results.
 	num is how many samples to try.
@@ -91,7 +91,7 @@ def compare_segments(dbconnect, base_dir='.', prototypes_path="./odo-digit-proto
 
 @cli
 @argh.named("analyze-segment")
-def do_analyze_segment(dbconnect, *segment_paths, base_dir='.', prototypes_path="./odo-digit-prototypes"):
+def do_analyze_segment(dbconnect, *segment_paths, base_dir='.', prototypes_path="./prototypes"):
 	"""Analyze individual segments and write them to the database"""
 	prototypes = load_prototypes(prototypes_path)
 	dbmanager = database.DBManager(dsn=dbconnect)
@@ -182,7 +182,7 @@ def main(
 	hours=2,
 	run_once=False,
 	overwrite=False,
-	prototypes_path="./odo-digit-prototypes",
+	prototypes_path="./prototypes",
 ):
 	CHECK_INTERVAL = 2