Last active
March 19, 2019 05:13
-
-
Save xxv/c7634ff3d4c0e21bebb23e56ca1ea810 to your computer and use it in GitHub Desktop.
Google Coral OpenCV camera capture demo
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
Usage: | |
python3 demo/classify_capture_opencv.py \ | |
--model test_data/inception_v4_299_quant_edgetpu.tflite \ | |
--label test_data/imagenet_labels.txt | |
""" | |
import argparse | |
import io | |
import time | |
import numpy as np | |
import cv2 | |
import edgetpu.classification.engine | |
def main(): | |
parser = argparse.ArgumentParser() | |
parser.add_argument( | |
'--model', help='File path of Tflite model.', required=True) | |
parser.add_argument( | |
'--label', help='File path of label file.', required=True) | |
args = parser.parse_args() | |
with open(args.label, 'r') as f: | |
pairs = (l.strip().split(maxsplit=1) for l in f.readlines()) | |
labels = dict((int(k), v) for k, v in pairs) | |
engine = edgetpu.classification.engine.ClassificationEngine(args.model) | |
try: | |
cap = cv2.VideoCapture(0) | |
font = cv2.FONT_HERSHEY_SIMPLEX | |
_, width, height, channels = engine.get_input_tensor_shape() | |
while True: | |
ret, frame = cap.read() | |
# Display the resulting frame | |
if cv2.waitKey(1) & 0xFF == ord('q'): | |
break | |
resized = cv2.resize(frame, (width, height)) | |
input = np.frombuffer(resized, dtype=np.uint8) | |
start_time = time.time() | |
results = engine.ClassifyWithInputTensor(input, top_k=1) | |
elapsed_time = time.time() - start_time | |
if results: | |
confidence = results[0][1] | |
label = labels[results[0][0]] | |
print("Elapsed time: {:0.02f}".format(elapsed_time * 1000)) | |
cv2.putText(frame, label, (0, 30), font, 1, (255, 255, 255), 2, cv2.LINE_AA) | |
cv2.putText(frame, "{:0.02f}".format(confidence), (0, 50), font, 1, (255, 255, 255), 2, cv2.LINE_AA) | |
cv2.imshow('frame', frame) | |
finally: | |
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