-
Notifications
You must be signed in to change notification settings - Fork 201
/
Copy pathdemo.py
156 lines (135 loc) · 6.87 KB
/
demo.py
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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
# 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.
import sys
import argparse
import numpy as np
import cv2 as cv
# Check OpenCV version
opencv_python_version = lambda str_version: tuple(map(int, (str_version.split("."))))
assert opencv_python_version(cv.__version__) >= opencv_python_version("4.10.0"), \
"Please install latest opencv-python for benchmark: python3 -m pip install --upgrade opencv-python"
from sface import SFace
sys.path.append('../face_detection_yunet')
from yunet import YuNet
# Valid combinations of backends and targets
backend_target_pairs = [
[cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_TARGET_CPU],
[cv.dnn.DNN_BACKEND_CUDA, cv.dnn.DNN_TARGET_CUDA],
[cv.dnn.DNN_BACKEND_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16],
[cv.dnn.DNN_BACKEND_TIMVX, cv.dnn.DNN_TARGET_NPU],
[cv.dnn.DNN_BACKEND_CANN, cv.dnn.DNN_TARGET_NPU]
]
parser = argparse.ArgumentParser(
description="SFace: Sigmoid-Constrained Hypersphere Loss for Robust Face Recognition (https://ieeexplore.ieee.org/document/9318547)")
parser.add_argument('--target', '-t', type=str,
help='Usage: Set path to the input image 1 (target face).')
parser.add_argument('--query', '-q', type=str,
help='Usage: Set path to the input image 2 (query).')
parser.add_argument('--model', '-m', type=str, default='face_recognition_sface_2021dec.onnx',
help='Usage: Set model path, defaults to face_recognition_sface_2021dec.onnx.')
parser.add_argument('--backend_target', '-bt', type=int, default=0,
help='''Choose one of the backend-target pair to run this demo:
{:d}: (default) OpenCV implementation + CPU,
{:d}: CUDA + GPU (CUDA),
{:d}: CUDA + GPU (CUDA FP16),
{:d}: TIM-VX + NPU,
{:d}: CANN + NPU
'''.format(*[x for x in range(len(backend_target_pairs))]))
parser.add_argument('--dis_type', type=int, choices=[0, 1], default=0,
help='Usage: Distance type. \'0\': cosine, \'1\': norm_l1. Defaults to \'0\'')
parser.add_argument('--save', '-s', action='store_true',
help='Usage: Specify to save file with results (i.e. bounding box, confidence level). Invalid in case of camera input.')
parser.add_argument('--vis', '-v', action='store_true',
help='Usage: Specify to open a new window to show results. Invalid in case of camera input.')
args = parser.parse_args()
def visualize(img1, faces1, img2, faces2, matches, scores, target_size=[512, 512]): # target_size: (h, w)
out1 = img1.copy()
out2 = img2.copy()
matched_box_color = (0, 255, 0) # BGR
mismatched_box_color = (0, 0, 255) # BGR
# Resize to 256x256 with the same aspect ratio
padded_out1 = np.zeros((target_size[0], target_size[1], 3)).astype(np.uint8)
h1, w1, _ = out1.shape
ratio1 = min(target_size[0] / out1.shape[0], target_size[1] / out1.shape[1])
new_h1 = int(h1 * ratio1)
new_w1 = int(w1 * ratio1)
resized_out1 = cv.resize(out1, (new_w1, new_h1), interpolation=cv.INTER_LINEAR).astype(np.float32)
top = max(0, target_size[0] - new_h1) // 2
bottom = top + new_h1
left = max(0, target_size[1] - new_w1) // 2
right = left + new_w1
padded_out1[top : bottom, left : right] = resized_out1
# Draw bbox
bbox1 = faces1[0][:4] * ratio1
x, y, w, h = bbox1.astype(np.int32)
cv.rectangle(padded_out1, (x + left, y + top), (x + left + w, y + top + h), matched_box_color, 2)
# Resize to 256x256 with the same aspect ratio
padded_out2 = np.zeros((target_size[0], target_size[1], 3)).astype(np.uint8)
h2, w2, _ = out2.shape
ratio2 = min(target_size[0] / out2.shape[0], target_size[1] / out2.shape[1])
new_h2 = int(h2 * ratio2)
new_w2 = int(w2 * ratio2)
resized_out2 = cv.resize(out2, (new_w2, new_h2), interpolation=cv.INTER_LINEAR).astype(np.float32)
top = max(0, target_size[0] - new_h2) // 2
bottom = top + new_h2
left = max(0, target_size[1] - new_w2) // 2
right = left + new_w2
padded_out2[top : bottom, left : right] = resized_out2
# Draw bbox
assert faces2.shape[0] == len(matches), "number of faces2 needs to match matches"
assert len(matches) == len(scores), "number of matches needs to match number of scores"
for index, match in enumerate(matches):
bbox2 = faces2[index][:4] * ratio2
x, y, w, h = bbox2.astype(np.int32)
box_color = matched_box_color if match else mismatched_box_color
cv.rectangle(padded_out2, (x + left, y + top), (x + left + w, y + top + h), box_color, 2)
score = scores[index]
text_color = matched_box_color if match else mismatched_box_color
cv.putText(padded_out2, "{:.2f}".format(score), (x + left, y + top - 5), cv.FONT_HERSHEY_DUPLEX, 0.4, text_color)
return np.concatenate([padded_out1, padded_out2], axis=1)
if __name__ == '__main__':
backend_id = backend_target_pairs[args.backend_target][0]
target_id = backend_target_pairs[args.backend_target][1]
# Instantiate SFace for face recognition
recognizer = SFace(modelPath=args.model,
disType=args.dis_type,
backendId=backend_id,
targetId=target_id)
# Instantiate YuNet for face detection
detector = YuNet(modelPath='../face_detection_yunet/face_detection_yunet_2023mar.onnx',
inputSize=[320, 320],
confThreshold=0.9,
nmsThreshold=0.3,
topK=5000,
backendId=backend_id,
targetId=target_id)
img1 = cv.imread(args.target)
img2 = cv.imread(args.query)
# Detect faces
detector.setInputSize([img1.shape[1], img1.shape[0]])
faces1 = detector.infer(img1)
assert faces1.shape[0] > 0, 'Cannot find a face in {}'.format(args.target)
detector.setInputSize([img2.shape[1], img2.shape[0]])
faces2 = detector.infer(img2)
assert faces2.shape[0] > 0, 'Cannot find a face in {}'.format(args.query)
# Match
scores = []
matches = []
for face in faces2:
result = recognizer.match(img1, faces1[0][:-1], img2, face[:-1])
scores.append(result[0])
matches.append(result[1])
# Draw results
image = visualize(img1, faces1, img2, faces2, matches, scores)
# Save results if save is true
if args.save:
print('Resutls saved to result.jpg\n')
cv.imwrite('result.jpg', image)
# Visualize results in a new window
if args.vis:
cv.namedWindow("SFace Demo", cv.WINDOW_AUTOSIZE)
cv.imshow("SFace Demo", image)
cv.waitKey(0)