File size: 1,938 Bytes
b03e0d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
# This file is part of OpenCV Zoo project.
# It is subject to the license terms in the LICENSE file found in the same directory.
#
# Copyright (C) 2021, Shenzhen Institute of Artificial Intelligence and Robotics for Society, all rights reserved.
# Third party copyrights are property of their respective owners.
from itertools import product
import numpy as np
import cv2 as cv
class YuNet:
def __init__(self, modelPath, inputSize=[320, 320], confThreshold=0.6, nmsThreshold=0.3, topK=5000, backendId=0, targetId=0):
self._modelPath = modelPath
self._inputSize = tuple(inputSize) # [w, h]
self._confThreshold = confThreshold
self._nmsThreshold = nmsThreshold
self._topK = topK
self._backendId = backendId
self._targetId = targetId
self._model = cv.FaceDetectorYN.create(
model=self._modelPath,
config="",
input_size=self._inputSize,
score_threshold=self._confThreshold,
nms_threshold=self._nmsThreshold,
top_k=self._topK,
backend_id=self._backendId,
target_id=self._targetId)
@property
def name(self):
return self.__class__.__name__
def setBackendAndTarget(self, backendId, targetId):
self._backendId = backendId
self._targetId = targetId
self._model = cv.FaceDetectorYN.create(
model=self._modelPath,
config="",
input_size=self._inputSize,
score_threshold=self._confThreshold,
nms_threshold=self._nmsThreshold,
top_k=self._topK,
backend_id=self._backendId,
target_id=self._targetId)
def setInputSize(self, input_size):
self._model.setInputSize(tuple(input_size))
def infer(self, image):
# Forward
faces = self._model.detect(image)
return np.array([]) if faces[1] is None else faces[1]
|