Skip to content

Commit 3c4aaac

Browse files
committed
tflite-multiprocessing on Pi
1 parent 07219b7 commit 3c4aaac

3 files changed

Lines changed: 98 additions & 0 deletions

File tree

File renamed without changes.
File renamed without changes.
Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
import os
2+
import cv2
3+
import multiprocessing as mp
4+
import numpy as np
5+
import time
6+
import psutil
7+
import tflite_runtime.interpreter as tflite
8+
9+
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
10+
MODEL_PATH = os.path.join(ROOT, "model_int8.tflite")
11+
12+
def camera_worker(frame_queue):
13+
cap = cv2.VideoCapture(0)
14+
while True:
15+
ret, frame = cap.read()
16+
if not ret:
17+
continue
18+
if not frame_queue.full():
19+
frame_queue.put(frame)
20+
21+
def classifier_worker(frame_queue, result_queue):
22+
interpreter = tflite.Interpreter(model_path=MODEL_PATH, num_threads=2)
23+
interpreter.allocate_tensors()
24+
25+
input_details = interpreter.get_input_details()
26+
output_details = interpreter.get_output_details()
27+
28+
input_shape = input_details[0]['shape']
29+
height, width = input_shape[1], input_shape[2]
30+
31+
while True:
32+
if not frame_queue.empty():
33+
frame = frame_queue.get()
34+
35+
resized_frame = cv2.resize(frame, (width, height))
36+
37+
input_data = np.expand_dims(resized_frame, axis=0)
38+
39+
if input_details[0]['dtype'] == np.float32:
40+
input_data = input_data.astype(np.float32) / 255.0
41+
else:
42+
input_data = input_data.astype(input_details[0]['dtype'])
43+
44+
interpreter.set_tensor(input_details[0]['index'], input_data)
45+
interpreter.invoke()
46+
47+
output_data = interpreter.get_tensor(output_details[0]['index'])
48+
49+
pred_idx = np.argmax(output_data)
50+
pred_text = f"Class: {pred_idx}"
51+
52+
result_queue.put((frame, pred_text))
53+
54+
def display_worker(result_queue):
55+
prev_frame_time = 0
56+
psutil.cpu_percent()
57+
58+
while True:
59+
if not result_queue.empty():
60+
frame, pred = result_queue.get()
61+
62+
new_frame_time = time.time()
63+
if prev_frame_time != 0:
64+
fps = 1 / (new_frame_time - prev_frame_time)
65+
else:
66+
fps = 0
67+
prev_frame_time = new_frame_time
68+
69+
cpu_usage = psutil.cpu_percent()
70+
71+
fps_text = f"FPS: {int(fps)}"
72+
cpu_text = f"CPU: {cpu_usage}%"
73+
74+
cv2.putText(frame, str(pred), (30,30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,255,0), 2)
75+
cv2.putText(frame, fps_text, (30,70), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,0,0), 2)
76+
cv2.putText(frame, cpu_text, (30,110), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
77+
78+
cv2.imshow("Multiprocess Classifier", frame)
79+
80+
if cv2.waitKey(1) == 27:
81+
break
82+
83+
if __name__=="__main__":
84+
frame_queue = mp.Queue(maxsize=5)
85+
result_queue = mp.Queue(maxsize=5)
86+
87+
p1 = mp.Process(target=camera_worker, args=(frame_queue,))
88+
p2 = mp.Process(target=classifier_worker, args=(frame_queue, result_queue))
89+
p3 = mp.Process(target=display_worker, args=(result_queue,))
90+
91+
p1.start()
92+
p2.start()
93+
p3.start()
94+
95+
p1.join()
96+
p2.join()
97+
p3.join()
98+

0 commit comments

Comments
 (0)