tensorflow lite raspberry pi object detection using usb webcam

1import time
2import re
3import numpy as np
4
5from tflite_runtime.interpreter import Interpreter
6import cv2
7
8CAMERA_WIDTH = 640
9CAMERA_HEIGHT = 480
10
11
12def load_labels(path):
13  """Loads the labels file. Supports files with or without index numbers."""
14  with open(path, 'r', encoding='utf-8') as f:
15    lines = f.readlines()
16    labels = {}
17    for row_number, content in enumerate(lines):
18      pair = re.split(r'[:\s]+', content.strip(), maxsplit=1)
19      if len(pair) == 2 and pair[0].strip().isdigit():
20        labels[int(pair[0])] = pair[1].strip()
21      else:
22        labels[row_number] = pair[0].strip()
23  return labels
24
25
26def set_input_tensor(interpreter, image):
27  """Sets the input tensor."""
28  tensor_index = interpreter.get_input_details()[0]['index']
29  input_tensor = interpreter.tensor(tensor_index)()[0]
30  
31  input_tensor[:, :] = cv2.resize(image,(300,300))
32  
33
34
35def get_output_tensor(interpreter, index):
36  """Returns the output tensor at the given index."""
37  output_details = interpreter.get_output_details()[index]
38  tensor = np.squeeze(interpreter.get_tensor(output_details['index']))
39  return tensor
40
41
42def detect_objects(interpreter, image, threshold):
43  """Returns a list of detection results, each a dictionary of object info."""
44  set_input_tensor(interpreter, image)
45  interpreter.invoke()
46
47  # Get all output details
48  boxes = get_output_tensor(interpreter, 0)
49  classes = get_output_tensor(interpreter, 1)
50  scores = get_output_tensor(interpreter, 2)
51  count = int(get_output_tensor(interpreter, 3))
52
53  results = []
54  for i in range(count):
55    if scores[i] >= threshold:
56      result = {
57          'bounding_box': boxes[i],
58          'class_id': classes[i],
59          'score': scores[i]
60      }
61      results.append(result)
62  return results
63
64
65def main():
66  labels = load_labels('coco_labels.txt')
67  interpreter = Interpreter('detect.tflite')
68  interpreter.allocate_tensors()
69  _, input_height, input_width, _ = interpreter.get_input_details()[0]['shape']
70  cam = cv2.VideoCapture(1)
71  try:
72      ret = True
73      while ret:
74          for i in range(10):
75              ret, frame = cam.read()
76          image = frame
77          
78          start_time = time.monotonic()
79          results = detect_objects(interpreter, image, 0.5)
80          elapsed_ms = (time.monotonic() - start_time) * 1000
81          for obj in results:
82              ymin, xmin, ymax, xmax = obj['bounding_box']
83              xmin = int(xmin * CAMERA_WIDTH)
84              xmax = int(xmax * CAMERA_WIDTH)
85              ymin = int(ymin * CAMERA_HEIGHT)
86              ymax = int(ymax * CAMERA_HEIGHT)
87              classid = labels[obj['class_id']]
88              score = obj['score']
89              y = ymin + 10
90              if ymin - 10 > 10:
91                  y = ymin - 10
92              cv2.rectangle(image, (xmin, ymin),(xmax, ymax), (255, 255, 255), 1)
93              cv2.putText(image,classid+" "+ str(round(elapsed_ms)),(xmin, y),
94                          cv2.FONT_HERSHEY_PLAIN,
95                          1.5, (0, 255, 0),  1,
96                          cv2.LINE_AA,False)
97        
98          cv2.imshow('r',image)
99          if cv2.waitKey(1) & 0xFF == ord('q'):
100              break
101          #print(results)
102  finally:
103      cam.release()
104      cv2.destroyAllWindows()
105
106    
107
108if __name__ == '__main__':
109  main()
110
111