JeVoisBase  1.18
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
yunet.py
Go to the documentation of this file.
1 # This file is part of OpenCV Zoo project.
2 # It is subject to the license terms in the LICENSE file found in the same directory.
3 #
4 # Copyright (C) 2021, Shenzhen Institute of Artificial Intelligence and Robotics for Society, all rights reserved.
5 # Third party copyrights are property of their respective owners.
6 
7 from itertools import product
8 
9 import numpy as np
10 import cv2 as cv
11 
12 class YuNet:
13  def __init__(self, modelPath, inputSize=[320, 320], confThreshold=0.6, nmsThreshold=0.3, topK=5000, backendId=0, targetId=0):
14  self._modelPath = modelPath
15  self._inputSize = tuple(inputSize) # [w, h]
16  self._confThreshold = confThreshold
17  self._nmsThreshold = nmsThreshold
18  self._topK = topK
19  self._backendId = backendId
20  self._targetId = targetId
21 
22  self._model = cv.FaceDetectorYN.create(
23  model=self._modelPath,
24  config="",
25  input_size=self._inputSize,
26  score_threshold=self._confThreshold,
27  nms_threshold=self._nmsThreshold,
28  top_k=self._topK,
29  backend_id=self._backendId,
30  target_id=self._targetId)
31 
32  @property
33  def name(self):
34  return self.__class__.__name__
35 
36  def setBackend(self, backendId):
37  self._backendId = backendId
38  self._model = cv.FaceDetectorYN.create(
39  model=self._modelPath,
40  config="",
41  input_size=self._inputSize,
42  score_threshold=self._confThreshold,
43  nms_threshold=self._nmsThreshold,
44  top_k=self._topK,
45  backend_id=self._backendId,
46  target_id=self._targetId)
47 
48  def setTarget(self, targetId):
49  self._targetId = targetId
50  self._model = cv.FaceDetectorYN.create(
51  model=self._modelPath,
52  config="",
53  input_size=self._inputSize,
54  score_threshold=self._confThreshold,
55  nms_threshold=self._nmsThreshold,
56  top_k=self._topK,
57  backend_id=self._backendId,
58  target_id=self._targetId)
59 
60  def setInputSize(self, input_size):
61  self._model.setInputSize(tuple(input_size))
62 
63  def infer(self, image):
64  # Forward
65  faces = self._model.detect(image)
66  return faces[1]
67 
yunet.YuNet._nmsThreshold
_nmsThreshold
Definition: yunet.py:17
yunet.YuNet.setBackend
def setBackend(self, backendId)
Definition: yunet.py:36
yunet.YuNet.infer
def infer(self, image)
Definition: yunet.py:63
yunet.YuNet.name
def name(self)
Definition: yunet.py:33
yunet.YuNet._confThreshold
_confThreshold
Definition: yunet.py:16
yunet.YuNet.__init__
def __init__(self, modelPath, inputSize=[320, 320], confThreshold=0.6, nmsThreshold=0.3, topK=5000, backendId=0, targetId=0)
Definition: yunet.py:13
yunet.YuNet._topK
_topK
Definition: yunet.py:18
yunet.YuNet
Definition: yunet.py:12
yunet.YuNet._backendId
_backendId
Definition: yunet.py:19
yunet.YuNet._modelPath
_modelPath
Definition: yunet.py:14
yunet.YuNet._model
_model
Definition: yunet.py:22
yunet.YuNet._inputSize
_inputSize
Definition: yunet.py:15
yunet.YuNet.setInputSize
def setInputSize(self, input_size)
Definition: yunet.py:60
yunet.YuNet._targetId
_targetId
Definition: yunet.py:20
yunet.YuNet.setTarget
def setTarget(self, targetId)
Definition: yunet.py:48