Display model title in overlay.
By default title is a model name or it can be customized via
command line.
Change-Id: Ib92e8365268b1377fe73d8d1ca6b89b5391b814e
diff --git a/edgetpuvision/classify.py b/edgetpuvision/classify.py
index 11f6d29..51119ad 100644
--- a/edgetpuvision/classify.py
+++ b/edgetpuvision/classify.py
@@ -14,8 +14,8 @@
from edgetpu.classification.engine import ClassificationEngine
from . import svg
+from . import utils
from .apps import run_app
-from .utils import load_labels, input_image_size, same_input_image_sizes, avg_fps_counter
CSS_STYLES = str(svg.CssStyle({'.back': svg.Style(fill='black',
@@ -94,17 +94,17 @@
acc = accumulator(size=args.window, top_k=args.top_k)
acc.send(None) # Initialize.
- fps_counter=avg_fps_counter(30)
+ fps_counter = utils.avg_fps_counter(30)
- engines = [ClassificationEngine(m) for m in args.model.split(',')]
- assert same_input_image_sizes(engines)
+ engines, titles = utils.make_engines(args.model, ClassificationEngine)
+ assert utils.same_input_image_sizes(engines)
engines = itertools.cycle(engines)
engine = next(engines)
- labels = load_labels(args.labels)
+ labels = utils.load_labels(args.labels)
draw_overlay = True
- yield input_image_size(engine)
+ yield utils.input_image_size(engine)
output = None
while True:
@@ -121,7 +121,8 @@
if args.print:
print_results(inference_rate, results)
- output = overlay(None, results, inference_time, inference_rate, layout)
+ title = titles[engine]
+ output = overlay(title, results, inference_time, inference_rate, layout)
else:
output = None
diff --git a/edgetpuvision/detect.py b/edgetpuvision/detect.py
index c688729..b50f795 100644
--- a/edgetpuvision/detect.py
+++ b/edgetpuvision/detect.py
@@ -20,8 +20,8 @@
from edgetpu.detection.engine import DetectionEngine
from . import svg
+from . import utils
from .apps import run_app
-from .utils import load_labels, input_image_size, same_input_image_sizes, avg_fps_counter
CSS_STYLES = str(svg.CssStyle({'.back': svg.Style(fill='black',
stroke='black',
@@ -127,20 +127,20 @@
print(' %d: %s, area=%.2f' % (i, obj, obj.bbox.area()))
def render_gen(args):
- fps_counter=avg_fps_counter(30)
+ fps_counter = utils.avg_fps_counter(30)
- engines = [DetectionEngine(m) for m in args.model.split(',')]
- assert same_input_image_sizes(engines)
+ engines, titles = utils.make_engines(args.model, DetectionEngine)
+ assert utils.same_input_image_sizes(engines)
engines = itertools.cycle(engines)
engine = next(engines)
- labels = load_labels(args.labels) if args.labels else None
+ labels = utils.load_labels(args.labels) if args.labels else None
filtered_labels = set(l.strip() for l in args.filter.split(',')) if args.filter else None
get_color = make_get_color(args.color, labels)
draw_overlay = True
- yield input_image_size(engine)
+ yield utils.input_image_size(engine)
output = None
while True:
@@ -149,7 +149,7 @@
inference_rate = next(fps_counter)
if draw_overlay:
start = time.monotonic()
- objs = engine.DetectWithInputTensor(tensor, threshold=args.threshold, top_k=args.top_k)
+ objs = engine .DetectWithInputTensor(tensor, threshold=args.threshold, top_k=args.top_k)
inference_time = time.monotonic() - start
objs = [convert(obj, labels) for obj in objs]
@@ -161,7 +161,8 @@
if args.print:
print_results(inference_rate, objs)
- output = overlay(None, objs, get_color, inference_time, inference_rate, layout)
+ title = titles[engine]
+ output = overlay(title, objs, get_color, inference_time, inference_rate, layout)
else:
output = None
diff --git a/edgetpuvision/utils.py b/edgetpuvision/utils.py
index 8de053e..c0e2e49 100644
--- a/edgetpuvision/utils.py
+++ b/edgetpuvision/utils.py
@@ -1,4 +1,5 @@
import collections
+import os
import re
import time
@@ -27,3 +28,15 @@
window.append(curr - prev)
prev = curr
yield len(window) / sum(window)
+
+def make_engines(models, engine_class):
+ engines, titles = [], {}
+ for model in models.split(','):
+ if '@' in model:
+ model_path, title = model.split('@')
+ else:
+ model_path, title = model, os.path.basename(os.path.normpath(model))
+ engine = engine_class(model_path)
+ engines.append(engine)
+ titles[engine] = title
+ return engines, titles
