ytfeng commited on
Commit
edcfc33
·
1 Parent(s): a524e9b

Add options for demo scripts to select backend & targets (#43)

Browse files

* add options for selecting backend & targets

* add eol

Files changed (2) hide show
  1. demo.py +19 -2
  2. yunet.py +2 -1
demo.py CHANGED
@@ -19,9 +19,23 @@ def str2bool(v):
19
  else:
20
  raise NotImplementedError
21
 
 
 
 
 
 
 
 
 
 
 
 
 
22
  parser = argparse.ArgumentParser(description='YuNet: A Fast and Accurate CNN-based Face Detector (https://github.com/ShiqiYu/libfacedetection).')
23
  parser.add_argument('--input', '-i', type=str, help='Path to the input image. Omit for using default camera.')
24
  parser.add_argument('--model', '-m', type=str, default='face_detection_yunet_2021dec.onnx', help='Path to the model.')
 
 
25
  parser.add_argument('--conf_threshold', type=float, default=0.9, help='Filter out faces of confidence < conf_threshold.')
26
  parser.add_argument('--nms_threshold', type=float, default=0.3, help='Suppress bounding boxes of iou >= nms_threshold.')
27
  parser.add_argument('--top_k', type=int, default=5000, help='Keep top_k bounding boxes before NMS.')
@@ -61,7 +75,9 @@ if __name__ == '__main__':
61
  inputSize=[320, 320],
62
  confThreshold=args.conf_threshold,
63
  nmsThreshold=args.nms_threshold,
64
- topK=args.top_k)
 
 
65
 
66
  # If input is an image
67
  if args.input is not None:
@@ -117,4 +133,5 @@ if __name__ == '__main__':
117
  # Visualize results in a new Window
118
  cv.imshow('YuNet Demo', frame)
119
 
120
- tm.reset()
 
 
19
  else:
20
  raise NotImplementedError
21
 
22
+ backends = [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_BACKEND_CUDA]
23
+ targets = [cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16]
24
+ help_msg_backends = "Choose one of the computation backends: {:d}: OpenCV implementation (default); {:d}: CUDA"
25
+ help_msg_targets = "Chose one of the target computation devices: {:d}: CPU (default); {:d}: CUDA; {:d}: CUDA fp16"
26
+ try:
27
+ backends += [cv.dnn.DNN_BACKEND_TIMVX]
28
+ targets += [cv.dnn.DNN_TARGET_NPU]
29
+ help_msg_backends += "; {:d}: TIMVX"
30
+ help_msg_targets += "; {:d}: NPU"
31
+ except:
32
+ print('This version of OpenCV does not support TIM-VX and NPU. Visit https://gist.github.com/fengyuentau/5a7a5ba36328f2b763aea026c43fa45f for more information.')
33
+
34
  parser = argparse.ArgumentParser(description='YuNet: A Fast and Accurate CNN-based Face Detector (https://github.com/ShiqiYu/libfacedetection).')
35
  parser.add_argument('--input', '-i', type=str, help='Path to the input image. Omit for using default camera.')
36
  parser.add_argument('--model', '-m', type=str, default='face_detection_yunet_2021dec.onnx', help='Path to the model.')
37
+ parser.add_argument('--backend', '-b', type=int, default=backends[0], help=help_msg_backends.format(*backends))
38
+ parser.add_argument('--target', '-t', type=int, default=targets[0], help=help_msg_targets.format(*targets))
39
  parser.add_argument('--conf_threshold', type=float, default=0.9, help='Filter out faces of confidence < conf_threshold.')
40
  parser.add_argument('--nms_threshold', type=float, default=0.3, help='Suppress bounding boxes of iou >= nms_threshold.')
41
  parser.add_argument('--top_k', type=int, default=5000, help='Keep top_k bounding boxes before NMS.')
 
75
  inputSize=[320, 320],
76
  confThreshold=args.conf_threshold,
77
  nmsThreshold=args.nms_threshold,
78
+ topK=args.top_k,
79
+ backendId=args.backend,
80
+ targetId=args.target)
81
 
82
  # If input is an image
83
  if args.input is not None:
 
133
  # Visualize results in a new Window
134
  cv.imshow('YuNet Demo', frame)
135
 
136
+ tm.reset()
137
+
yunet.py CHANGED
@@ -63,4 +63,5 @@ class YuNet:
63
  def infer(self, image):
64
  # Forward
65
  faces = self._model.detect(image)
66
- return faces[1]
 
 
63
  def infer(self, image):
64
  # Forward
65
  faces = self._model.detect(image)
66
+ return faces[1]
67
+