3
3
from urllib .request import urlretrieve
4
4
5
5
import cv2
6
- from motpy import Detection , MultiObjectTracker , NpImage
6
+ from motpy import Detection , MultiObjectTracker , NpImage , Box
7
7
from motpy .core import setup_logger
8
8
from motpy .detector import BaseObjectDetector
9
9
from motpy .testing_viz import draw_detection , draw_track
15
15
16
16
"""
17
17
18
- logger = setup_logger (__name__ , is_main = True )
18
+ logger = setup_logger (__name__ , 'DEBUG' , is_main = True )
19
19
20
20
21
21
WEIGHTS_URL = 'https://github.com/opencv/opencv_3rdparty/raw/dnn_samples_face_detector_20170830/res10_300x300_ssd_iter_140000.caffemodel'
@@ -30,7 +30,7 @@ def __init__(self,
30
30
weights_path : str = WEIGHTS_PATH ,
31
31
config_url : str = CONFIG_URL ,
32
32
config_path : str = CONFIG_PATH ,
33
- conf_threshold : float = 0.5 ):
33
+ conf_threshold : float = 0.5 ) -> None :
34
34
super (FaceDetector , self ).__init__ ()
35
35
36
36
if not os .path .isfile (weights_path ) or not os .path .isfile (config_path ):
@@ -43,23 +43,23 @@ def __init__(self,
43
43
# specify detector hparams
44
44
self .conf_threshold = conf_threshold
45
45
46
- def process_image (self , image : NpImage ) -> Sequence [NpImage ]:
46
+ def process_image (self , image : NpImage ) -> Sequence [Detection ]:
47
47
blob = cv2 .dnn .blobFromImage (image , 1.0 , (300 , 300 ), [104 , 117 , 123 ], False , False )
48
48
self .net .setInput (blob )
49
49
detections = self .net .forward ()
50
50
51
51
# convert output from OpenCV detector to tracker expected format [xmin, ymin, xmax, ymax]
52
- bboxes = []
52
+ out_detections = []
53
53
for i in range (detections .shape [2 ]):
54
54
confidence = detections [0 , 0 , i , 2 ]
55
55
if confidence > self .conf_threshold :
56
56
xmin = int (detections [0 , 0 , i , 3 ] * image .shape [1 ])
57
57
ymin = int (detections [0 , 0 , i , 4 ] * image .shape [0 ])
58
58
xmax = int (detections [0 , 0 , i , 5 ] * image .shape [1 ])
59
59
ymax = int (detections [0 , 0 , i , 6 ] * image .shape [0 ])
60
- bboxes .append ([xmin , ymin , xmax , ymax ])
61
-
62
- return bboxes
60
+ out_detections .append (Detection ( box = [xmin , ymin , xmax , ymax ], score = confidence ) )
61
+
62
+ return out_detections
63
63
64
64
65
65
def run ():
@@ -84,8 +84,7 @@ def run():
84
84
frame = cv2 .resize (frame , dsize = None , fx = 0.5 , fy = 0.5 )
85
85
86
86
# run face detector on current frame
87
- bboxes = face_detector .process_image (frame )
88
- detections = [Detection (box = bbox ) for bbox in bboxes ]
87
+ detections = face_detector .process_image (frame )
89
88
logger .debug (f'detections: { detections } ' )
90
89
91
90
tracker .step (detections )
0 commit comments