mirror of
https://github.com/deepinsight/insightface.git
synced 2025-12-30 16:12:27 +00:00
98 lines
4.1 KiB
Python
98 lines
4.1 KiB
Python
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()
|