5
5
import darknet
6
6
import argparse
7
7
import threading
8
+ import multiprocessing as mp
8
9
import queue
9
10
10
11
@@ -26,6 +27,8 @@ def parser():
26
27
help = "path to data file" )
27
28
parser .add_argument ("--thresh" , type = float , default = .25 ,
28
29
help = "remove detections with confidence below this value" )
30
+ parser .add_argument ("--multiprocess" , action = "store_true" ,
31
+ help = "use processes instead of threads" )
29
32
return parser .parse_args ()
30
33
31
34
@@ -107,7 +110,7 @@ def convert4cropping(image, bbox, preproc_h, preproc_w):
107
110
108
111
def video_capture (stop_flag , input_path , raw_frame_queue , preprocessed_frame_queue , preproc_h , preproc_w ):
109
112
cap = cv2 .VideoCapture (input_path )
110
- while cap .isOpened () and not stop_flag .is_set ():
113
+ while cap .isOpened () and stop_flag .empty ():
111
114
ret , frame = cap .read ()
112
115
if not ret :
113
116
break
@@ -118,30 +121,38 @@ def video_capture(stop_flag, input_path, raw_frame_queue, preprocessed_frame_que
118
121
img_for_detect = darknet .make_image (preproc_w , preproc_h , 3 )
119
122
darknet .copy_image_from_bytes (img_for_detect , frame_resized .tobytes ())
120
123
preprocessed_frame_queue .put (img_for_detect )
121
- stop_flag .set ( )
124
+ stop_flag .put ( None )
122
125
cap .release ()
126
+ print ("video_capture end:" , os .getpid ())
123
127
124
128
125
129
def inference (stop_flag , preprocessed_frame_queue , detections_queue , fps_queue ,
126
- network , class_names , threshold ):
127
- while not stop_flag .is_set ():
130
+ config_file , data_file , weights_file , batch_size , threshold , ext_output ):
131
+ network , class_names , _ = darknet .load_network (
132
+ config_file ,
133
+ data_file ,
134
+ weights_file ,
135
+ batch_size = batch_size )
136
+ while stop_flag .empty ():
128
137
darknet_image = preprocessed_frame_queue .get ()
129
138
prev_time = time .time ()
130
139
detections = darknet .detect_image (network , class_names , darknet_image , thresh = threshold )
131
140
fps = 1 / (time .time () - prev_time )
132
141
detections_queue .put (detections )
133
142
fps_queue .put (int (fps ))
134
143
print ("FPS: {:.2f}" .format (fps ))
135
- darknet .print_detections (detections , args . ext_output )
144
+ # darknet.print_detections(detections, ext_output)
136
145
darknet .free_image (darknet_image )
146
+ darknet .free_network_ptr (network )
147
+ print ("inference end:" , os .getpid ())
137
148
138
149
139
- def drawing (stop_flag , input_video_fps , queues , preproc_h , preproc_w , vid_h , vid_w ):
150
+ def drawing (stop_flag , input_video_fps , queues , preproc_h , preproc_w , vid_h , vid_w , out_filename , dont_show , class_colors ):
140
151
random .seed (3 ) # deterministic bbox colors
141
152
raw_frame_queue , preprocessed_frame_queue , detections_queue , fps_queue = queues
142
- video = set_saved_video (args . out_filename , (vid_w , vid_h ), input_video_fps )
153
+ video = set_saved_video (out_filename , (vid_w , vid_h ), input_video_fps )
143
154
fps = 1
144
- while not stop_flag .is_set ():
155
+ while stop_flag .empty ():
145
156
frame = raw_frame_queue .get ()
146
157
detections = detections_queue .get ()
147
158
fps = fps_queue .get ()
@@ -151,13 +162,13 @@ def drawing(stop_flag, input_video_fps, queues, preproc_h, preproc_w, vid_h, vid
151
162
bbox_adjusted = convert2original (frame , bbox , preproc_h , preproc_w )
152
163
detections_adjusted .append ((str (label ), confidence , bbox_adjusted ))
153
164
image = darknet .draw_boxes (detections_adjusted , frame , class_colors )
154
- if not args . dont_show :
165
+ if not dont_show :
155
166
cv2 .imshow ("Inference" , image )
156
- if args . out_filename is not None :
167
+ if out_filename is not None :
157
168
video .write (image )
158
169
if cv2 .waitKey (fps ) == 27 :
159
170
break
160
- stop_flag .set ( )
171
+ stop_flag .put ( None )
161
172
video .release ()
162
173
cv2 .destroyAllWindows ()
163
174
timeout = 1 / (fps if fps > 0 else 0.5 )
@@ -166,18 +177,22 @@ def drawing(stop_flag, input_video_fps, queues, preproc_h, preproc_w, vid_h, vid
166
177
q .get (block = True , timeout = timeout )
167
178
except queue .Empty :
168
179
pass
180
+ print ("drawing end:" , os .getpid ())
169
181
170
182
171
183
if __name__ == "__main__" :
172
184
args = parser ()
173
185
check_arguments_errors (args )
174
- network , class_names , class_colors = darknet .load_network (
186
+ batch_size = 1
187
+ network , class_names , class_colors = darknet .load_network ( # Load network twice :(
175
188
args .config_file ,
176
189
args .data_file ,
177
190
args .weights ,
178
- batch_size = 1 )
191
+ batch_size = batch_size )
179
192
darknet_width = darknet .network_width (network )
180
193
darknet_height = darknet .network_height (network )
194
+ darknet .free_network_ptr (network )
195
+ del network
181
196
input_path = str2int (args .input )
182
197
cap = cv2 .VideoCapture (input_path ) # Open video twice :(
183
198
video_width = int (cap .get (cv2 .CAP_PROP_FRAME_WIDTH ))
@@ -186,9 +201,14 @@ def drawing(stop_flag, input_video_fps, queues, preproc_h, preproc_w, vid_h, vid
186
201
cap .release ()
187
202
del cap
188
203
189
- ExecUnit = threading .Thread
190
- Queue = queue .Queue
191
- stop_flag = threading .Event ()
204
+ if args .multiprocess :
205
+ ExecUnit = mp .Process
206
+ Queue = mp .Queue
207
+ else :
208
+ ExecUnit = threading .Thread
209
+ Queue = queue .Queue
210
+
211
+ stop_flag = Queue ()
192
212
193
213
raw_frame_queue = Queue ()
194
214
preprocessed_frame_queue = Queue (maxsize = 1 )
@@ -199,15 +219,19 @@ def drawing(stop_flag, input_video_fps, queues, preproc_h, preproc_w, vid_h, vid
199
219
ExecUnit (target = video_capture , args = (stop_flag , input_path , raw_frame_queue , preprocessed_frame_queue ,
200
220
darknet_height , darknet_width )),
201
221
ExecUnit (target = inference , args = (stop_flag , preprocessed_frame_queue , detections_queue , fps_queue ,
202
- network , class_names , args .thresh )),
222
+ args .config_file , args .data_file , args .weights , batch_size , args .thresh ,
223
+ args .ext_output )),
203
224
ExecUnit (target = drawing , args = (stop_flag , video_fps ,
204
225
(raw_frame_queue , preprocessed_frame_queue , detections_queue , fps_queue ),
205
- darknet_height , darknet_width , video_height , video_width )),
226
+ darknet_height , darknet_width , video_height , video_width ,
227
+ args .out_filename , args .dont_show , class_colors )),
206
228
)
207
229
for exec_unit in exec_units :
208
230
exec_unit .start ()
231
+ print ("------- EXEC UNIT:" , ExecUnit )
209
232
for exec_unit in exec_units :
210
233
exec_unit .join ()
234
+ print ("------- EXEC UNIT:" , ExecUnit )
211
235
212
236
print ("\n Done." )
213
237
0 commit comments