-
Notifications
You must be signed in to change notification settings - Fork 8
/
rk3588.py
128 lines (113 loc) · 4.35 KB
/
rk3588.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import json
from multiprocessing import Process, Queue
from pathlib import Path
from rknnlite.api import RKNNLite
from base.camera import Cam
from base.inference import Yolov5
from base.post_process import post_process
CONFIG_FILE = str(Path(__file__).parent.parent.absolute()) + "/config.json"
with open(CONFIG_FILE, 'r') as config_file:
cfg = json.load(config_file)
class Rk3588():
"""Class for object detection on RK3588/RK3588S
Attributes
---------------------------------------------------------------------------
Queues
-----------------------------------
_q_pre : multiprocessing.Queue
Queue for sending raw frames, resized frames and frames ids from
camera reading process to inference process
_q_outs : multiprocessing.Queue
Queue for sending inference results, raw frames and frames ids from
inference process to post_process process
_q_post : multiprocessing.Queue
Queue for sending raw frames, frames with bboxes, numpy array with
detections and frames ids from post_process process to ouput
-----------------------------------
Camera
-----------------------------------
cam : camera.Cam
Camera object for creating recording, showing process
-----------------------------------
Inference
-----------------------------------
_yolov5 : inference.Yolov5 or inference.VariableYolov5
Yolov5 object for creating inference processes
-----------------------------------
Processes
-----------------------------------
_rec : multiprocessing.Process
Process for recording frames
_inf : multiprocessing.Process
Process for inferencing frames (recomended amount is 3 and should equal
post_process processes)
_post : multiprocessing.Process
Process for post processing frames (recomended amount is 3 and should
equal inference processes)
-----------------------------------
---------------------------------------------------------------------------
Methods
---------------------------------------------------------------------------
start() : None
Starts all processes (recording process, inference process(es),
post_process process(es))
show() : None
Create cv2 window with inferenced frames (frames with bboxes on them)
get_data() : tuple(np.ndarray, np.ndarray, np.ndarray, int) | None
Returns raw frames, frames with bboxes, numpy array with detections
and frames ids
---------------------------------------------------------------------------
"""
def __init__(self):
self._q_pre = Queue(maxsize=cfg["inference"]["buf_size"])
self._q_outs = Queue(maxsize=cfg["inference"]["buf_size"])
self._q_post = Queue(maxsize=cfg["inference"]["buf_size"])
self._cam = Cam(
source = cfg["camera"]["source"],
q_in = self._q_post,
q_out = self._q_pre
)
self._cores=[
RKNNLite.NPU_CORE_0,
RKNNLite.NPU_CORE_1,
RKNNLite.NPU_CORE_2
]
self._yolov5 = [
Yolov5(
proc = i,
q_in = self._q_pre,
q_out = self._q_outs,
core=self._cores[i%3]
) for i in range(cfg["inference"]["inf_proc"])
]
self._rec = Process(
target = self._cam.record,
daemon=True
)
self._inf = [
Process(
target = self._yolov5[i].inference,
daemon = True
) for i in range(len(self._yolov5))
]
self._post = [
Process(
target = post_process,
kwargs = {
"q_in" : self._q_outs,
"q_out" : self._q_post
},
daemon=True
) for i in range(cfg["inference"]["post_proc"])
]
def start(self):
self._rec.start()
for inference in self._inf: inference.start()
for post_process in self._post: post_process.start()
def show(self, start_time):
self._cam.show(start_time)
def get_data(self):
if self._q_post.empty():
return None
raw_frame, inferenced_frame, detections, frame_id = self._q_post.get()
return(raw_frame, inferenced_frame, detections, frame_id)