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

* add options for selecting backend & targets

* add eol

demo.py

@@ -19,9 +19,23 @@ def str2bool(v):
     else:
         raise NotImplementedError
 
+backends = [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_BACKEND_CUDA]
+targets = [cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16]
+help_msg_backends = "Choose one of the computation backends: {:d}: OpenCV implementation (default); {:d}: CUDA"
+help_msg_targets = "Chose one of the target computation devices: {:d}: CPU (default); {:d}: CUDA; {:d}: CUDA fp16"
+try:
+    backends += [cv.dnn.DNN_BACKEND_TIMVX]
+    targets += [cv.dnn.DNN_TARGET_NPU]
+    help_msg_backends += "; {:d}: TIMVX"
+    help_msg_targets += "; {:d}: NPU"
+except:
+    print('This version of OpenCV does not support TIM-VX and NPU. Visit https://gist.github.com/fengyuentau/5a7a5ba36328f2b763aea026c43fa45f for more information.')
+
 parser = argparse.ArgumentParser(description='PPHumanSeg (https://github.com/PaddlePaddle/PaddleSeg/tree/release/2.2/contrib/PP-HumanSeg)')
 parser.add_argument('--input', '-i', type=str, help='Path to the input image. Omit for using default camera.')
 parser.add_argument('--model', '-m', type=str, default='human_segmentation_pphumanseg_2021oct.onnx', help='Path to the model.')
+parser.add_argument('--backend', '-b', type=int, default=backends[0], help=help_msg_backends.format(*backends))
+parser.add_argument('--target', '-t', type=int, default=targets[0], help=help_msg_targets.format(*targets))
 parser.add_argument('--save', '-s', type=str, default=False, help='Set true to save results. This flag is invalid when using camera.')
 parser.add_argument('--vis', '-v', type=str2bool, default=True, help='Set true to open a window for result visualization. This flag is invalid when using camera.')
 args = parser.parse_args()
@@ -84,7 +98,7 @@ def visualize(image, result, weight=0.6, fps=None):
 
 if __name__ == '__main__':
     # Instantiate PPHumanSeg
-    model = PPHumanSeg(modelPath=args.model)
+    model = PPHumanSeg(modelPath=args.model, backendId=args.backend, targetId=args.target)
 
     if args.input is not None:
         # Read image and resize to 192x192
@@ -138,4 +152,5 @@ if __name__ == '__main__':
             # Visualize results in a new window
             cv.imshow('PPHumanSeg Demo', frame)
 
-            tm.reset()
+            tm.reset()
+

pphumanseg.py

@@ -8,9 +8,14 @@ import numpy as np
 import cv2 as cv
 
 class PPHumanSeg:
-    def __init__(self, modelPath):
+    def __init__(self, modelPath, backendId=0, targetId=0):
         self._modelPath = modelPath
+        self._backendId = backendId
+        self._targetId = targetId
+
         self._model = cv.dnn.readNet(self._modelPath)
+        self._model.setPreferableBackend(self._backendId)
+        self._model.setPreferableTarget(self._targetId)
 
         self._inputNames = ''
         self._outputNames = ['save_infer_model/scale_0.tmp_1']
@@ -23,10 +28,12 @@ class PPHumanSeg:
         return self.__class__.__name__
 
     def setBackend(self, backend_id):
-        self._model.setPreferableBackend(backend_id)
+        self._backendId = backend_id
+        self._model.setPreferableBackend(self._backendId)
 
     def setTarget(self, target_id):
-        self._model.setPreferableTarget(target_id)
+        self._targetId = target_id
+        self._model.setPreferableTarget(self._targetId)
 
     def _preprocess(self, image):
         image = image.astype(np.float32, copy=False) / 255.0
@@ -52,4 +59,5 @@ class PPHumanSeg:
 
     def _postprocess(self, outputBlob):
         result = np.argmax(outputBlob[0], axis=1).astype(np.uint8)
-        return result
+        return result
+