mirror of
https://github.com/yakhyo/uniface.git
synced 2025-12-30 09:02:25 +00:00
77f14a616a
- add dynamic onnx provider selection for m1/m2/m3/m4 macs - replace mkdocs with simple markdown files - fix model download and scrfd detection issues - update ci/cd workflows
87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from uniface.constants import RetinaFaceWeights
|
|
from uniface.detection import RetinaFace
|
|
|
|
|
|
@pytest.fixture
|
|
def retinaface_model():
|
|
"""
|
|
Fixture to initialize the RetinaFace model for testing.
|
|
"""
|
|
return RetinaFace(
|
|
model_name=RetinaFaceWeights.MNET_V2,
|
|
conf_thresh=0.5,
|
|
pre_nms_topk=5000,
|
|
nms_thresh=0.4,
|
|
post_nms_topk=750,
|
|
)
|
|
|
|
|
|
def test_model_initialization(retinaface_model):
|
|
"""
|
|
Test that the RetinaFace model initializes correctly.
|
|
"""
|
|
assert retinaface_model is not None, "Model initialization failed."
|
|
|
|
|
|
def test_inference_on_640x640_image(retinaface_model):
|
|
"""
|
|
Test inference on a 640x640 BGR image.
|
|
"""
|
|
# Generate a mock 640x640 BGR image
|
|
mock_image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
|
|
|
|
# Run inference - returns list of dictionaries
|
|
faces = retinaface_model.detect(mock_image)
|
|
|
|
# Check output type
|
|
assert isinstance(faces, list), "Detections should be a list."
|
|
|
|
# Check that each face has the expected structure
|
|
for face in faces:
|
|
assert isinstance(face, dict), "Each detection should be a dictionary."
|
|
assert "bbox" in face, "Each detection should have a 'bbox' key."
|
|
assert "confidence" in face, "Each detection should have a 'confidence' key."
|
|
assert "landmarks" in face, "Each detection should have a 'landmarks' key."
|
|
|
|
# Check bbox format
|
|
bbox = face["bbox"]
|
|
assert len(bbox) == 4, "BBox should have 4 values (x1, y1, x2, y2)."
|
|
|
|
# Check landmarks format
|
|
landmarks = face["landmarks"]
|
|
assert len(landmarks) == 5, "Should have 5 landmark points."
|
|
assert all(len(pt) == 2 for pt in landmarks), "Each landmark should be (x, y)."
|
|
|
|
|
|
def test_confidence_threshold(retinaface_model):
|
|
"""
|
|
Test that detections respect the confidence threshold.
|
|
"""
|
|
# Generate a mock 640x640 BGR image
|
|
mock_image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
|
|
|
|
# Run inference
|
|
faces = retinaface_model.detect(mock_image)
|
|
|
|
# Ensure all detections have confidence scores above the threshold
|
|
for face in faces:
|
|
confidence = face["confidence"]
|
|
assert confidence >= 0.5, f"Detection has confidence {confidence} below threshold 0.5"
|
|
|
|
|
|
def test_no_faces_detected(retinaface_model):
|
|
"""
|
|
Test inference on an image without detectable faces.
|
|
"""
|
|
# Generate an empty (black) 640x640 image
|
|
empty_image = np.zeros((640, 640, 3), dtype=np.uint8)
|
|
|
|
# Run inference
|
|
faces = retinaface_model.detect(empty_image)
|
|
|
|
# Ensure no detections are found
|
|
assert len(faces) == 0, "Should detect no faces in a blank image."
|