Skip to content

Instantly share code, notes, and snippets.

@Namburger
Last active January 20, 2020 17:06
Show Gist options
  • Save Namburger/b64502ca017399c9427ccd909a5b914e to your computer and use it in GitHub Desktop.
Save Namburger/b64502ca017399c9427ccd909a5b914e to your computer and use it in GitHub Desktop.
An example of doing object detection with open cv on the dev board
# Copyright 2019 Google LLC
#
# Changes made by Nam Vu 01/20/2020
#
# 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.
r"""A demo for object detection.
Model and Label File taken from here:
https://github.com/google-coral/edgetpu/tree/master/test_data
$ python3 object_detection.py --model mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite --label coco_labels.txt
"""
import argparse
import platform
import subprocess
from edgetpu.detection.engine import DetectionEngine
from edgetpu.utils import dataset_utils
import cv2
import numpy as np
from PIL import Image
from PIL import ImageDraw
def do_detection(engine, img, k=3):
# Run inference.
ans = engine.detect_with_image(
img,
threshold=0.05,
keep_aspect_ratio=False,
relative_coord=False,
top_k=k)
return ans
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--model',
help='Path of the detection model, it must be a SSD model with postprocessing operator.',
required=True)
parser.add_argument('--label', help='Path of the labels file.')
parser.add_argument('--top_k', help='Top number to results to return.')
args = parser.parse_args()
# Initialize engine.
engine = DetectionEngine(args.model)
labels = dataset_utils.read_label_file(args.label) if args.label else None
top_k = args.top_k if args.top_k else 3
# capturing videos with cv2
cap = cv2.VideoCapture(0)
# looping through all frames
frame_num = 0
while True:
print('<<<----------------------------------------->>>')
print('Frame number: ', frame_num)
ret, cv2img = cap.read()
cv2img = cv2.cvtColor(cv2img, cv2.COLOR_BGR2RGB)
img = Image.fromarray(cv2img)
draw = ImageDraw.Draw(img)
ans = do_detection(engine, img)
# Save result.
if ans:
for obj in ans:
print('-----------------------------------------')
if labels:
print(labels[obj.label_id])
print('score = ', obj.score)
box = obj.bounding_box.flatten().tolist()
print('box=', box)
draw.rectangle(box, outline='red')
else:
print('No object detected!')
new_cv2img = np.array(img.convert('RGB'))
new_cv2img = new_cv2img[:, :, ::-1].copy() # RGB to BGR
cv2.imshow('Object Detection', new_cv2img)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
frame_num += 1
print('\n')
cap.release()
cv2.destroyAllWindows()
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment