Add convenient object representation for detection.
Change-Id: I38b7552ad76e71c13a68250458bf40075fcaa2bf
diff --git a/edgetpuvision/detect.py b/edgetpuvision/detect.py
index ed102e9..6a38c3e 100644
--- a/edgetpuvision/detect.py
+++ b/edgetpuvision/detect.py
@@ -12,26 +12,36 @@
# --labels ${TEST_DATA}/coco_labels.txt
import argparse
+import collections
import itertools
import time
from edgetpu.detection.engine import DetectionEngine
-
from . import overlays
from .utils import load_labels, input_image_size, same_input_image_sizes, avg_fps_counter
from .gstreamer import Display, run_gen
-def area(obj):
- x0, y0, x1, y1 = rect = obj.bounding_box.flatten().tolist()
- return (x1 - x0) * (y1 - y0)
+BBox = collections.namedtuple('BBox', ('x', 'y', 'w', 'h'))
+BBox.area = lambda self: self.w * self.h
+BBox.scale = lambda self, sx, sy: BBox(x=self.x * sx, y=self.y * sy,
+ w=self.w * sx, h=self.h * sy)
+BBox.__str__ = lambda self: 'BBox(x=%.2f y=%.2f w=%.2f h=%.2f)' % self
-def print_results(inference_rate, objs, labels):
+Object = collections.namedtuple('Object', ('id', 'label', 'score', 'bbox'))
+Object.__str__ = lambda self: 'Object(id=%d, label=%s, score=%.2f, %s)' % self
+
+def convert(obj, labels):
+ x0, y0, x1, y1 = obj.bounding_box.flatten().tolist()
+ return Object(id=obj.label_id,
+ label=labels[obj.label_id] if labels else None,
+ score=obj.score,
+ bbox=BBox(x=x0, y=y0, w=x1 - x0, h=y1 - y0))
+
+def print_results(inference_rate, objs):
print('\nInference (rate=%.2f fps):' % inference_rate)
for i, obj in enumerate(objs):
- label = labels[obj.label_id] if labels else str(obj.label_id)
- x = (i, label) + tuple(obj.bounding_box.flatten()) + (area(obj),)
- print(' %d: label=%s, bbox=(%.2f %.2f %.2f %.2f), bbox_area=%.2f' % x)
+ print(' %d: %s, area=%.2f' % (i, obj, obj.bbox.area()))
def render_gen(args):
fps_counter=avg_fps_counter(30)
@@ -56,16 +66,17 @@
start = time.monotonic()
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]
if labels and filtered_labels:
- objs = [obj for obj in objs if labels[obj.label_id] in filtered_labels]
+ objs = [obj for obj in objs if obj.label in filtered_labels]
- objs = [obj for obj in objs if args.min_area <= area(obj) <= args.max_area]
+ objs = [obj for obj in objs if args.min_area <= obj.bbox.area() <= args.max_area]
if args.print:
- print_results(inference_rate, objs, labels)
+ print_results(inference_rate, objs)
- output = overlays.detection(objs, labels, inference_time, inference_rate, layout)
+ output = overlays.detection(objs, inference_time, inference_rate, layout)
else:
output = None
diff --git a/edgetpuvision/overlays.py b/edgetpuvision/overlays.py
index cfdb458..bcc8fd6 100644
--- a/edgetpuvision/overlays.py
+++ b/edgetpuvision/overlays.py
@@ -4,14 +4,6 @@
'.shd': svg.Style(fill='black', fill_opacity=0.6),
'rect': svg.Style(fill='green', fill_opacity=0.3, stroke='white')}))
-
-def _normalize_rect(rect, size):
- width, height = size
- x0, y0, x1, y1 = rect
- return int(x0 * width), int(y0 * height), \
- int((x1 - x0) * width), int((y1 - y0) * height)
-
-
def classification(results, inference_time, inference_rate, layout):
x0, y0, w, h = layout.window
@@ -31,7 +23,7 @@
doc += svg.normal_text(lines, x=x0 + 10, y=y0 + 10, font_size_em=1.1)
return str(doc)
-def detection(objs, labels, inference_time, inference_rate, layout):
+def detection(objs, inference_time, inference_rate, layout):
x0, y0, w, h = layout.window
defs = svg.Defs()
@@ -47,12 +39,12 @@
for obj in objs:
percent = int(100 * obj.score)
- if labels:
- caption = '%d%% %s' % (percent, labels[obj.label_id])
+ if obj.label:
+ caption = '%d%% %s' % (percent, obj.label)
else:
caption = '%d%%' % percent
- x, y, w, h = _normalize_rect(obj.bounding_box.flatten().tolist(), layout.size)
+ x, y, w, h = obj.bbox.scale(*layout.size)
doc += svg.normal_text(caption, x, y - 5)
doc += svg.Rect(x=x, y=y, width=w, height=h, rx=2, ry=2)
