Spaces:
Sleeping
Sleeping
File size: 3,594 Bytes
c741d15 5699031 c741d15 |
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 |
import cv2 as cv
import numpy as np
import gradio as gr
from huggingface_hub import hf_hub_download
from yunet import YuNet
from ediffiqa import eDifFIQA
# Download face detection model (YuNet)
model_path_yunet = hf_hub_download(
repo_id="opencv/face_detection_yunet",
filename="face_detection_yunet_2023mar.onnx"
)
# Download face quality assessment model (eDifFIQA Tiny)
model_path_quality = hf_hub_download(
repo_id="opencv/face_image_quality_assessment_ediffiqa",
filename="ediffiqa_tiny_jun2024.onnx"
)
# Backend and target
backend_id = cv.dnn.DNN_BACKEND_OPENCV
target_id = cv.dnn.DNN_TARGET_CPU
# Initialize YuNet for face detection
face_detector = YuNet(
modelPath=model_path_yunet,
inputSize=[320, 320],
confThreshold=0.9,
nmsThreshold=0.3,
topK=5000,
backendId=backend_id,
targetId=target_id
)
# Initialize eDifFIQA for quality assessment
quality_model = eDifFIQA(
modelPath=model_path_quality,
inputSize=[112, 112]
)
quality_model.setBackendAndTarget(
backendId=backend_id,
targetId=target_id
)
REFERENCE_FACIAL_POINTS = np.array([
[38.2946 , 51.6963 ],
[73.5318 , 51.5014 ],
[56.0252 , 71.7366 ],
[41.5493 , 92.3655 ],
[70.729904, 92.2041 ]
], dtype=np.float32)
def align_image(image, detection_data):
src_pts = np.float32(detection_data[0][4:-1]).reshape(5, 2)
tfm, _ = cv.estimateAffinePartial2D(src_pts, REFERENCE_FACIAL_POINTS, method=cv.LMEDS)
face_img = cv.warpAffine(image, tfm, (112, 112))
return face_img
def assess_face_quality(input_image):
bgr_image = cv.cvtColor(input_image, cv.COLOR_RGB2BGR)
h, w, _ = bgr_image.shape
face_detector.setInputSize([w, h])
detections = face_detector.infer(bgr_image)
if detections is None or len(detections) == 0:
return "No face detected.", input_image
aligned_face = align_image(bgr_image, detections)
score = np.squeeze(quality_model.infer(aligned_face)).item()
output_image = aligned_face.copy()
cv.putText(output_image, f"{score:.3f}", (0, 20), cv.FONT_HERSHEY_DUPLEX, 0.8, (0, 0, 255), 2)
output_image = cv.cvtColor(output_image, cv.COLOR_BGR2RGB)
return f"Quality Score: {score:.3f}", output_image
# Gradio Interface
with gr.Blocks(css='''.example * {
font-style: italic;
font-size: 18px !important;
color: #0ea5e9 !important;
}''') as demo:
gr.Markdown("### Face Image Quality Assessment (eDifFIQA + YuNet)")
gr.Markdown("Upload a face image. The app detects and aligns the face, then evaluates image quality using the eDifFIQA model.")
with gr.Row():
input_image = gr.Image(type="numpy", label="Upload Face Image")
with gr.Column():
quality_score = gr.Text(label="Quality Score")
aligned_face = gr.Image(type="numpy", label="Aligned Face with Score")
# Clear output when new image is uploaded
input_image.change(fn=lambda: ("", None), outputs=[quality_score, aligned_face])
with gr.Row():
submit_btn = gr.Button("Submit", variant="primary")
clear_btn = gr.Button("Clear")
submit_btn.click(fn=assess_face_quality, inputs=input_image, outputs=[quality_score, aligned_face])
clear_btn.click(fn=lambda: (None, "", None), outputs=[input_image, quality_score, aligned_face])
gr.Markdown("Click on any example to try it.", elem_classes=["example"])
gr.Examples(
examples=[
["examples/lena.jpg"],
["examples/gray_face.png"]
],
inputs=input_image
)
if __name__ == "__main__":
demo.launch()
|