5
5
import torch
6
6
import torch .nn as nn
7
7
import torch .nn .functional as F
8
- from torch .multiprocessing import Process , Queue
8
+ import torch .multiprocessing as mp
9
9
import time
10
10
from PIL import Image
11
11
import numpy as np
12
12
import cv2
13
13
14
- import lib .transform_cv2 as T
14
+ import lib .data . transform_cv2 as T
15
15
from lib .models import model_factory
16
16
from configs import set_cfg_from_file
17
17
@@ -40,7 +40,7 @@ def get_model():
40
40
41
41
42
42
# fetch frames
43
- def get_func (inpth , in_q ):
43
+ def get_func (inpth , in_q , done ):
44
44
cap = cv2 .VideoCapture (args .input )
45
45
width = cap .get (cv2 .CAP_PROP_FRAME_WIDTH ) # type is float
46
46
height = cap .get (cv2 .CAP_PROP_FRAME_HEIGHT ) # type is float
@@ -59,7 +59,8 @@ def get_func(inpth, in_q):
59
59
in_q .put (frame )
60
60
61
61
in_q .put ('quit' )
62
- while not in_q .empty (): continue
62
+ done .wait ()
63
+
63
64
cap .release ()
64
65
time .sleep (1 )
65
66
print ('input queue done' )
@@ -105,14 +106,15 @@ def infer_batch(frames):
105
106
106
107
107
108
if __name__ == '__main__' :
108
- torch . multiprocessing .set_start_method ('spawn' )
109
+ mp .set_start_method ('spawn' )
109
110
110
- in_q = Queue (1024 )
111
- out_q = Queue (1024 )
111
+ in_q = mp .Queue (1024 )
112
+ out_q = mp .Queue (1024 )
113
+ done = mp .Event ()
112
114
113
- in_worker = Process (target = get_func ,
114
- args = (args .input , in_q ))
115
- out_worker = Process (target = save_func ,
115
+ in_worker = mp . Process (target = get_func ,
116
+ args = (args .input , in_q , done ))
117
+ out_worker = mp . Process (target = save_func ,
116
118
args = (args .input , args .output , out_q ))
117
119
118
120
in_worker .start ()
@@ -133,6 +135,7 @@ def infer_batch(frames):
133
135
infer_batch (frames )
134
136
135
137
out_q .put ('quit' )
138
+ done .set ()
136
139
137
140
out_worker .join ()
138
141
in_worker .join ()
0 commit comments