From 4fb9f101b4f8c2fbf3fa8feae6d2d22fd3262678 Mon Sep 17 00:00:00 2001 From: Mats Alritzson Date: Mon, 29 Jul 2019 09:04:31 +0200 Subject: [PATCH 1/4] Added detector --- pygame/detect.py | 102 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 pygame/detect.py diff --git a/pygame/detect.py b/pygame/detect.py new file mode 100644 index 0000000..a7bff08 --- /dev/null +++ b/pygame/detect.py @@ -0,0 +1,102 @@ +# 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 to classify Pygame camera stream.""" +import argparse +import os +import io +import time +import re +from collections import deque +import numpy as np +import pygame +import pygame.camera +from pygame.locals import * + +from edgetpu.detection.engine import DetectionEngine + +def load_labels(path): + p = re.compile(r'\s*(\d+)(.+)') + with open(path, 'r', encoding='utf-8') as f: + lines = (p.match(line).groups() for line in f.readlines()) + return {int(num): text.strip() for num, text in lines} + +def main(): + cam_w, cam_h = 640, 480 + default_model_dir = "../all_models" + default_model = 'mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite' + default_labels = 'coco_labels.txt' + parser = argparse.ArgumentParser() + parser.add_argument('--model', help='.tflite model path', + default=os.path.join(default_model_dir,default_model)) + parser.add_argument('--labels', help='label file path', + default=os.path.join(default_model_dir, default_labels)) + parser.add_argument('--top_k', type=int, default=5, + help='number of classes with highest score to display') + parser.add_argument('--threshold', type=float, default=0.5, + help='class score threshold') + args = parser.parse_args() + + with open(args.labels, 'r') as f: + pairs = (l.strip().split(maxsplit=1) for l in f.readlines()) + labels = dict((int(k), v) for k, v in pairs) + + print("Loading %s with %s labels."%(args.model, args.labels)) + engine = DetectionEngine(args.model) + labels = load_labels(args.labels) + + pygame.init() + pygame.font.init() + font = pygame.font.SysFont("Arial", 20) + + pygame.camera.init() + camlist = pygame.camera.list_cameras() + + _, w, h, _ = engine.get_input_tensor_shape() + camera = pygame.camera.Camera(camlist[0], (cam_w, cam_h)) + display = pygame.display.set_mode((cam_w, cam_h), 0) + + red = pygame.Color(255, 0, 0) + + camera.start() + try: + last_time = time.monotonic() + while True: + mysurface = camera.get_image() + imagen = pygame.transform.scale(mysurface, (w, h)) + input = np.frombuffer(imagen.get_buffer(), dtype=np.uint8) + start_time = time.monotonic() + results = engine.DetectWithInputTensor(input, threshold=args.threshold, top_k=args.top_k) + stop_time = time.monotonic() + inference_ms = (stop_time - start_time)*1000.0 + fps_ms = 1.0 / (stop_time - last_time) + last_time = stop_time + annotate_text = "Inference: %5.2fms FPS: %3.1f" % (inference_ms, fps_ms) + for result in results: + x0, y0, x1, y1 = result.bounding_box.flatten().tolist() + rect = pygame.Rect(x0 * cam_w, y0 * cam_h, (x1 - x0) * cam_w, (y1 - y0) * cam_h) + pygame.draw.rect(mysurface, red, rect, 1) + label = "%.0f%% %s" % (100*result.score, labels[result.label_id]) + text = font.render(label, True, red) + mysurface.blit(text, (x0 * cam_w , y0 * cam_h)) + text = font.render(annotate_text, True, red) + mysurface.blit(text, (0, 0)) + display.blit(mysurface, (0, 0)) + pygame.display.flip() + finally: + camera.stop() + + +if __name__ == '__main__': + main() From 63c3859176caa0511c01f85350cb2deada325f76 Mon Sep 17 00:00:00 2001 From: "mats.alritzson@volvocars.com" Date: Mon, 29 Jul 2019 09:12:20 +0200 Subject: [PATCH 2/4] Added some documentation --- pygame/README.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/pygame/README.md b/pygame/README.md index 59d3b97..5a41172 100644 --- a/pygame/README.md +++ b/pygame/README.md @@ -1,7 +1,6 @@ -This folder contains some simple camera classification examples using pygame +This folder contains some simple camera classification and detection examples using pygame If you dont have pygame installed you can install it by: - ``` pip3 install pygame ``` @@ -9,14 +8,18 @@ pip3 install pygame To run the demo execture the following command, which will use the default model ```mobilenet_v2_1.0_224_quant_edgetpu.tflite``` - +You run the classifier with: ``` python3 classify_capture.py ``` You can change the model and the labels file using flags: - ``` python3 classify_capture.py --model ../all_models/inception_v3_299_quant_edgetpu.tflite - ``` + +You run the detector with: +``` +python3 detect.py +``` + From b39702247385fd0c9eee71d717178dd444963c49 Mon Sep 17 00:00:00 2001 From: enmasse Date: Mon, 29 Jul 2019 11:14:08 +0200 Subject: [PATCH 3/4] Updated comment in detect.py --- pygame/detect.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pygame/detect.py b/pygame/detect.py index a7bff08..8af7a3d 100644 --- a/pygame/detect.py +++ b/pygame/detect.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""A demo to classify Pygame camera stream.""" +"""A demo to run the dector in a Pygame camera stream.""" import argparse import os import io From f06e9ec82a4529edded98faef782a47723720934 Mon Sep 17 00:00:00 2001 From: enmasse Date: Mon, 29 Jul 2019 11:15:30 +0200 Subject: [PATCH 4/4] Duh! --- pygame/detect.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pygame/detect.py b/pygame/detect.py index 8af7a3d..5349846 100644 --- a/pygame/detect.py +++ b/pygame/detect.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""A demo to run the dector in a Pygame camera stream.""" +"""A demo to run the detector in a Pygame camera stream.""" import argparse import os import io