blob: 165806f3443613bbe9a0447b0d7c3cb20add50eb [file] [log] [blame]
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""A demo which runs object classification on camera frames.
export TEST_DATA=/usr/lib/python3/dist-packages/edgetpu/test_data
python3 -m edgetpuvision.classify \
--model ${TEST_DATA}/mobilenet_v2_1.0_224_inat_bird_quant.tflite \
--labels ${TEST_DATA}/inat_bird_labels.txt
"""
import argparse
import collections
import itertools
import time
from pycoral.adapters import classify
from pycoral.utils import edgetpu
from . import svg
from . import utils
from .apps import run_app
CSS_STYLES = str(svg.CssStyle({'.back': svg.Style(fill='black',
stroke='black',
stroke_width='0.5em')}))
def size_em(length):
return '%sem' % str(0.6 * (length + 1))
def overlay(title, results, inference_time, inference_rate, layout):
x0, y0, width, height = layout.window
font_size = 0.03 * height
defs = svg.Defs()
defs += CSS_STYLES
doc = svg.Svg(width=width, height=height,
viewBox='%s %s %s %s' % layout.window,
font_size=font_size, font_family='monospace', font_weight=500)
doc += defs
ox1, ox2 = x0 + 20, x0 + width - 20
oy1, oy2 = y0 + 20 + font_size, y0 + height - 20
# Classes
lines = ['%s (%.2f)' % pair for pair in results]
for i, line in enumerate(lines):
y = oy2 - i * 1.7 * font_size
doc += svg.Rect(x=0, y=0, width=size_em(len(line)), height='1em',
transform='translate(%s, %s) scale(-1,-1)' % (ox2, y),
_class='back')
doc += svg.Text(line, text_anchor='end', x=ox2, y=y, fill='white')
# Title
if title:
doc += svg.Rect(x=0, y=0, width=size_em(len(title)), height='1em',
transform='translate(%s, %s) scale(1,-1)' % (ox1, oy1), _class='back')
doc += svg.Text(title, x=ox1, y=oy1, fill='white')
# Info
lines = [
'Inference time: %.2f ms (%.2f fps)' % (inference_time * 1000, 1.0 / inference_time)
]
for i, line in enumerate(reversed(lines)):
y = oy2 - i * 1.7 * font_size
doc += svg.Rect(x=0, y=0, width=size_em(len(line)), height='1em',
transform='translate(%s, %s) scale(1,-1)' % (ox1, y), _class='back')
doc += svg.Text(line, x=ox1, y=y, fill='white')
return str(doc)
def top_results(window, top_k):
total_scores = collections.defaultdict(lambda: 0.0)
for results in window:
for label, score in results:
total_scores[label] += score
return sorted(total_scores.items(), key=lambda kv: kv[1], reverse=True)[:top_k]
def accumulator(size, top_k):
window = collections.deque(maxlen=size)
window.append((yield []))
while True:
window.append((yield top_results(window, top_k)))
def print_results(inference_rate, results):
print('\nInference (rate=%.2f fps):' % inference_rate)
print(results)
for label, score in results:
print(' %s, score=%.2f' % (label, score))
def render_gen(args):
acc = accumulator(size=args.window, top_k=args.top_k)
acc.send(None) # Initialize.
fps_counter = utils.avg_fps_counter(30)
interpreters, titles = utils.make_interpreters(args.model)
assert utils.same_input_image_sizes(interpreters)
interpreters = itertools.cycle(interpreters)
interpreter = next(interpreters)
labels = utils.load_labels(args.labels)
draw_overlay = True
yield utils.input_image_size(interpreter)
output = None
while True:
tensor, layout, command = (yield output)
inference_rate = next(fps_counter)
if draw_overlay:
start = time.monotonic()
edgetpu.run_inference(interpreter, tensor)
inference_time = time.monotonic() - start
classes = classify.get_classes(interpreter, top_k=args.top_k,
score_threshold=args.threshold)
results = [(labels[class_id], score) for class_id, score in classes]
results = acc.send(results)
if args.print:
print_results(inference_rate, results)
title = titles[interpreter]
output = overlay(title, results, inference_time, inference_rate, layout)
else:
output = None
if command == 'o':
draw_overlay = not draw_overlay
elif command == 'n':
interpreter = next(interpreters)
def add_render_gen_args(parser):
parser.add_argument('--model', required=True,
help='.tflite model path')
parser.add_argument('--labels', required=True,
help='label file path')
parser.add_argument('--window', type=int, default=10,
help='number of frames to accumulate inference results')
parser.add_argument('--top_k', type=int, default=3,
help='number of classes with highest score to display')
parser.add_argument('--threshold', type=float, default=0.1,
help='class score threshold')
parser.add_argument('--print', default=False, action='store_true',
help='Print inference results')
def main():
run_app(add_render_gen_args, render_gen)
if __name__ == '__main__':
main()