diff --git a/deepface/models/facial_recognition/Facenet.py b/deepface/models/facial_recognition/Facenet.py index 799fe2d..e619392 100644 --- a/deepface/models/facial_recognition/Facenet.py +++ b/deepface/models/facial_recognition/Facenet.py @@ -84,7 +84,7 @@ class FaceNet512dONNXClient(FacialRecognition): input_name = self.model.get_inputs()[0].name output_name = self.model.get_outputs()[0].name result = self.model.run([output_name], {input_name: img}) - return result[0] + return result[0][0] def scaling(x, scale): @@ -1734,7 +1734,7 @@ def load_facenet512d_model( def load_facenet512d_onnx_model( - url="https://github.com/serengil/deepface_models/releases/download/v1.0/facenet512_weights.h5", + url="https://github.com/ShivamSinghal1/deepface/releases/download/v1/facenet512_fp32.onnx", ) -> Model: """ Download Facenet512d ONNX model weights and load @@ -1743,14 +1743,14 @@ def load_facenet512d_onnx_model( """ try: import onnxruntime as ort - import torch except ModuleNotFoundError as e: raise ImportError( "FaceNet512dONNX is an optional model, ensure the library is installed. " "Please install using 'pip install onnxruntime' or 'pip install onnxruntime-gpu' to use gpu" ) from e - if torch.cuda.is_available(): + if ort.get_device() == "GPU": + logger.info(f"using onnx GPU for inference") providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] else: providers = ['CPUExecutionProvider'] diff --git a/requirements_additional.txt b/requirements_additional.txt index ea76fde..100d887 100644 --- a/requirements_additional.txt +++ b/requirements_additional.txt @@ -3,4 +3,5 @@ mediapipe>=0.8.7.3 dlib>=19.20.0 ultralytics>=8.0.122 facenet-pytorch>=2.5.3 -torch>=2.1.2 \ No newline at end of file +torch>=2.1.2 +onnxruntime>=1.19.0 \ No newline at end of file diff --git a/tests/test_verify.py b/tests/test_verify.py index c1a80e4..47ccb4f 100644 --- a/tests/test_verify.py +++ b/tests/test_verify.py @@ -8,7 +8,7 @@ from deepface.commons.logger import Logger logger = Logger() -models = ["VGG-Face", "Facenet", "Facenet512", "ArcFace", "GhostFaceNet"] +models = ["VGG-Face", "Facenet", "Facenet512", "ArcFace", "GhostFaceNet", "Facenet512ONNX"] metrics = ["cosine", "euclidean", "euclidean_l2"] detectors = ["opencv", "mtcnn"]