-
Notifications
You must be signed in to change notification settings - Fork 5.6k
/
Copy pathsample_face_recognition.py
97 lines (83 loc) · 4.08 KB
/
sample_face_recognition.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
import os
import cv2
import inspireface as isf
import click
@click.command()
@click.argument('test_data_folder')
def case_face_recognition(test_data_folder):
"""
Launches the face recognition system, inserts face features into a database, and performs searches.
Args:
resource_path (str): Path to the resource directory for face recognition algorithms.
test_data_folder (str): Path to the test data containing images for insertion and recognition tests.
"""
# If you need to switch from the default Pikachu model to another model like Megatron, you can use reload
ret = isf.reload("Megatron")
assert ret, "Launch failure. Please ensure the resource path is correct."
# Enable face recognition features.
opt = isf.HF_ENABLE_FACE_RECOGNITION
session = isf.InspireFaceSession(opt, isf.HF_DETECT_MODE_ALWAYS_DETECT)
# Configure the feature management system.
feature_hub_config = isf.FeatureHubConfiguration(
primary_key_mode=isf.HF_PK_AUTO_INCREMENT,
enable_persistence=False,
persistence_db_path="",
search_threshold=0.48,
search_mode=isf.HF_SEARCH_MODE_EAGER,
)
ret = isf.feature_hub_enable(feature_hub_config)
assert ret, "Failed to enable FeatureHub."
# Insert face features from 'bulk' directory.
bulk_path = os.path.join(test_data_folder, "bulk")
assert os.path.exists(bulk_path), "Bulk directory does not exist."
insert_images = [os.path.join(bulk_path, path) for path in os.listdir(bulk_path) if path.endswith(".jpg")]
for idx, image_path in enumerate(insert_images):
name = os.path.basename(image_path).replace(".jpg", "")
image = cv2.imread(image_path)
assert image is not None, f"Failed to load image {image_path}"
faces = session.face_detection(image)
if faces:
face = faces[0] # Assume the most prominent face is what we want.
feature = session.face_feature_extract(image, face)
identity = isf.FaceIdentity(feature, id=idx)
ret, alloc_id = isf.feature_hub_face_insert(identity)
assert ret, "Failed to insert face."
count = isf.feature_hub_get_face_count()
print(f"Number of faces inserted: {count}")
# Process faces from 'RD' directory and insert them.
RD = os.path.join(test_data_folder, "RD")
assert os.path.exists(RD), "RD directory does not exist."
RD_images = [os.path.join(RD, path) for path in os.listdir(RD) if path.endswith(".jpeg")]
for idx, image_path in enumerate(RD_images[:-1]):
name = os.path.basename(image_path).replace(".jpeg", "")
image = cv2.imread(image_path)
assert image is not None, f"Failed to load image {image_path}"
faces = session.face_detection(image)
if faces:
face = faces[0]
feature = session.face_feature_extract(image, face)
identity = isf.FaceIdentity(feature, id=idx+count+1)
ret, alloc_id = isf.feature_hub_face_insert(identity)
assert ret, "Failed to insert face."
count = isf.feature_hub_get_face_count()
print(f"Total number of faces after insertion: {count}")
# Search for a similar face using the last image in RD directory.
remain = cv2.imread(RD_images[-1])
assert remain is not None, f"Failed to load image {RD_images[-1]}"
faces = session.face_detection(remain)
assert faces, "No faces detected."
face = faces[0]
feature = session.face_feature_extract(remain, face)
search = isf.feature_hub_face_search(feature)
if search.similar_identity.id != -1:
print(f"Found similar identity with ID: {search.similar_identity.id}, Confidence: {search.confidence:.2f}")
else:
print("No similar identity found.")
# Display top-k similar face identities.
print("Top-k similar identities:")
search_top_k = isf.feature_hub_face_search_top_k(feature, 10)
for idx, (conf, _id) in enumerate(search_top_k):
identity = isf.feature_hub_get_face_identity(_id)
print(f"Top-{idx + 1}: ID: {_id}, Confidence: {conf:.2f}")
if __name__ == '__main__':
case_face_recognition()