Advertisement
jarekmor

model_tflite_edge

Oct 12th, 2022
821
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.59 KB | None | 0 0
  1. import re
  2. import os
  3. import cv2
  4. from pycoral.utils.dataset import read_label_file
  5. from pycoral.utils.edgetpu import make_interpreter
  6. from pycoral.adapters import common
  7. from pycoral.adapters import classify
  8.  
  9.  
  10. # the TFLite converted to be used with edgetpu
  11. modelPath = './efficientdet-lite-ha_edgetpu.tflite'
  12.  
  13. # The path to labels.txt that was downloaded with your model
  14. labelPath = './ha-labels.txt'
  15.  
  16. # This function takes in a TFLite Interptere and Image, and returns classifications
  17. def classifyImage(interpreter, image):
  18.     size = common.input_size(interpreter)
  19.     common.set_input(interpreter, cv2.resize(image, size, fx=0, fy=0,
  20.                                              interpolation=cv2.INTER_CUBIC))
  21.     interpreter.invoke()
  22.     return classify.get_classes(interpreter)
  23.  
  24. def main():
  25.     # Load your model onto the TF Lite Interpreter
  26.     interpreter = make_interpreter(modelPath)
  27.     interpreter.allocate_tensors()
  28.     labels = read_label_file(labelPath)
  29.  
  30.     cap = cv2.VideoCapture('rtsp://192.168.1.212:8020/h264_pcm.sdp')
  31.     while cap.isOpened():
  32.         ret, frame = cap.read()
  33.         if not ret:
  34.             break
  35.  
  36.         # Flip image so it matches the training input
  37.         frame = cv2.flip(frame, 1)
  38.  
  39.         # Classify and display image
  40.         results = classifyImage(interpreter, frame)
  41. #        cv2.imshow('frame', frame)
  42.         print(f'Label: {labels[results[0].id]}, Score: {results[0].score}')
  43. #        if cv2.waitKey(1) & 0xFF == ord('q'):
  44. #            break
  45.  
  46.     cap.release()
  47.     cv2.destroyAllWindows()
  48.  
  49. if __name__ == '__main__':
  50.     main()
  51.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement