mirror of
https://github.com/yakhyo/uniface.git
synced 2025-12-30 00:52:25 +00:00
79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
import pytest
|
|
import numpy as np
|
|
from uniface import RetinaFace
|
|
|
|
|
|
@pytest.fixture
|
|
def retinaface_model():
|
|
"""
|
|
Fixture to initialize the RetinaFace model for testing.
|
|
"""
|
|
return RetinaFace(
|
|
model="retinaface_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
|
|
detections, landmarks = retinaface_model.detect(mock_image)
|
|
|
|
# Check output types
|
|
assert isinstance(detections, np.ndarray), "Detections should be a numpy array."
|
|
assert isinstance(landmarks, np.ndarray), "Landmarks should be a numpy array."
|
|
|
|
# Check that detections have the expected shape
|
|
if detections.size > 0: # If faces are detected
|
|
assert detections.shape[1] == 5, "Each detection should have 5 values (x1, y1, x2, y2, score)."
|
|
|
|
# Check landmarks shape
|
|
if landmarks.size > 0:
|
|
assert landmarks.shape[1:] == (5, 2), "Landmarks should have shape (N, 5, 2)."
|
|
|
|
|
|
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
|
|
detections, _ = retinaface_model.detect(mock_image)
|
|
|
|
# Ensure all detections have confidence scores above the threshold
|
|
if detections.size > 0: # If faces are detected
|
|
confidence_scores = detections[:, 4]
|
|
assert (confidence_scores >= 0.5).all(), "Some detections have confidence below the threshold."
|
|
|
|
|
|
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
|
|
detections, landmarks = retinaface_model.detect(empty_image)
|
|
|
|
# Ensure no detections or landmarks are found
|
|
assert detections.size == 0, "Detections should be empty for a blank image."
|
|
assert landmarks.size == 0, "Landmarks should be empty for a blank image."
|