Bug fix: use white color when labels file is not specified.
Change-Id: I672543ee27751e223d02bcf1e885d8849fc2b7fb
diff --git a/edgetpuvision/detect.py b/edgetpuvision/detect.py
index 42c5efb..1742da4 100644
--- a/edgetpuvision/detect.py
+++ b/edgetpuvision/detect.py
@@ -41,13 +41,20 @@
def color(i, total):
return tuple(int(255.0 * c) for c in colorsys.hsv_to_rgb(i / total, 1.0, 1.0))
-def gen_colors_unique(keys):
+def make_palette(keys):
return {key : svg.rgb(color(i, len(keys))) for i, key in enumerate(keys)}
-def gen_colors_single(keys, color):
- return {key : color for key in keys}
+def make_get_color(color, labels):
+ if color:
+ return lambda obj_id: color
-def overlay(objs, colors, inference_time, inference_rate, layout):
+ if labels:
+ palette = make_palette(labels.keys())
+ return lambda obj_id: palette[obj_id]
+
+ return lambda obj_id: 'white'
+
+def overlay(objs, get_color, inference_time, inference_rate, layout):
x0, y0, width, height = layout.window
defs = svg.Defs()
@@ -71,7 +78,7 @@
doc += svg.Rect(x=x + 1, y=y + 1, width=w, height=h, rx=2, ry=2,
_class='bbox', style='stroke:black')
doc += svg.Rect(x=x, y=y, width=w, height=h, rx=2, ry=2,
- _class='bbox', style='stroke:%s' % colors[obj.id])
+ _class='bbox', style='stroke:%s' % get_color(obj.id))
ox, oy = x0 + 20, y0 + height - 20
@@ -111,10 +118,8 @@
labels = 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
- if args.color:
- colors = gen_colors_single(labels.keys(), args.color)
- else:
- colors = gen_colors_unique(labels.keys())
+ get_color = make_get_color(args.color, labels)
+
draw_overlay = True
yield input_image_size(engine)
@@ -138,7 +143,7 @@
if args.print:
print_results(inference_rate, objs)
- output = overlay(objs, colors, inference_time, inference_rate, layout)
+ output = overlay(objs, get_color, inference_time, inference_rate, layout)
else:
output = None
