Yuantao Feng
Update to OpenCV APIs (YuNet -> FaceDetectorYN, SFace -> FaceRecognizerSF) (#6)
3af1dea
| # 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) | |
| def name(self): | |
| return self.__class__.__name__ | |
| def setBackend(self, backendId): | |
| self._backendId = backendId | |
| 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 setTarget(self, targetId): | |
| 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 faces[1] |