55 lines
2.0 KiB
Python
55 lines
2.0 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
输入：aligned cropped face
|
|||
|
|
输出：face feature
|
|||
|
|
注意： 提供计算两个feature的相似度的函数
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import cv2
|
|||
|
|
import onnxruntime as ort
|
|||
|
|
import numpy as np
|
|||
|
|
|
|||
|
|
# 设置ONNX Runtime的日志级别为ERROR
|
|||
|
|
ort.set_default_logger_severity(3) # 3表示ERROR级别
|
|||
|
|
|
|||
|
|
class FaceRecoger():
|
|||
|
|
def __init__(self, onnx_path, num_threads=1) -> None:
|
|||
|
|
session_options = ort.SessionOptions()
|
|||
|
|
session_options.intra_op_num_threads = num_threads
|
|||
|
|
# 初始化 InferenceSession 时传入 SessionOptions 对象
|
|||
|
|
self.ort_session = ort.InferenceSession(onnx_path, session_options=session_options)
|
|||
|
|
|
|||
|
|
output_node_dims = [out.shape for out in self.ort_session.get_outputs()]
|
|||
|
|
self.len_feat = output_node_dims[0][1] # feature 的长度为...
|
|||
|
|
|
|||
|
|
def inference(self, crop_img): # crop_img = cv2.imread(img_path) bgr
|
|||
|
|
input_feed = {}
|
|||
|
|
if crop_img.shape[:2] != (248,248): # 这里还有另一种方式 ,[4:252, 4:252,...]
|
|||
|
|
crop_img = cv2.resize(crop_img,(248,248))
|
|||
|
|
crop_img = crop_img[...,::-1]
|
|||
|
|
input_data = crop_img.transpose((2, 0, 1))
|
|||
|
|
input_feed['_input_123'] = input_data.reshape((1, 3, 248, 248)).astype(np.float32)
|
|||
|
|
pred_result = self.ort_session.run([], input_feed=input_feed)
|
|||
|
|
temp_result = np.sqrt(pred_result[0])
|
|||
|
|
norm = temp_result / np.linalg.norm(temp_result, axis=1)
|
|||
|
|
return norm.flatten() # return normalize feature
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def compute_sim(feat1,feat2):
|
|||
|
|
feat1, feat2 = feat1.flatten(), feat2.flatten()
|
|||
|
|
assert feat1.shape == feat2.shape
|
|||
|
|
sim = np.sum(feat1 * feat2)
|
|||
|
|
return sim
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
fr = FaceRecoger(onnx_path = "./checkpoints/face_recognizer.onnx", num_threads=1)
|
|||
|
|
import sys
|
|||
|
|
imgpath1 = sys.argv[1]
|
|||
|
|
imgpath2 = sys.argv[2]
|
|||
|
|
img1,img2 = cv2.imread(imgpath1),cv2.imread(imgpath2)
|
|||
|
|
feat1 = fr.inference(img1)
|
|||
|
|
feat2 = fr.inference(img2)
|
|||
|
|
|
|||
|
|
print("sim: ", FaceRecoger.compute_sim(feat1, feat2))
|