Abhishek Gola
Adding face detection to opencv space
d828c23
raw
history blame
2.01 kB
import cv2 as cv
import numpy as np
import gradio as gr
from yunet import YuNet
from huggingface_hub import hf_hub_download
# Download ONNX model from Hugging Face
model_path = hf_hub_download(repo_id="opencv/face_detection_yunet", filename="face_detection_yunet_2023mar.onnx")
# Initialize YuNet model
model = YuNet(
modelPath=model_path,
inputSize=[320, 320],
confThreshold=0.9,
nmsThreshold=0.3,
topK=5000,
backendId=cv.dnn.DNN_BACKEND_OPENCV,
targetId=cv.dnn.DNN_TARGET_CPU
)
def visualize(image, results, box_color=(0, 255, 0), text_color=(0, 0, 255)):
output = image.copy()
landmark_color = [
(255, 0, 0), # right eye
( 0, 0, 255), # left eye
( 0, 255, 0), # nose tip
(255, 0, 255), # right mouth corner
( 0, 255, 255) # left mouth corner
]
for det in results:
bbox = det[0:4].astype(np.int32)
cv.rectangle(output, (bbox[0], bbox[1]), (bbox[0]+bbox[2], bbox[1]+bbox[3]), box_color, 2)
conf = det[-1]
cv.putText(output, '{:.2f}'.format(conf), (bbox[0], bbox[1] + 12), cv.FONT_HERSHEY_SIMPLEX, 0.5, text_color)
landmarks = det[4:14].astype(np.int32).reshape((5, 2))
for idx, landmark in enumerate(landmarks):
cv.circle(output, tuple(landmark), 2, landmark_color[idx], 2)
return output
def detect_faces(input_image):
h, w, _ = input_image.shape
model.setInputSize([w, h])
results = model.infer(input_image)
if results is None or len(results) == 0:
return input_image
return visualize(input_image, results)
# Gradio Interface
demo = gr.Interface(
fn=detect_faces,
inputs=gr.Image(type="numpy", label="Upload Image"),
outputs=gr.Image(type="numpy", label="Detected Faces"),
title="Face Detection YuNet (OpenCV DNN)",
allow_flagging="never",
description="Upload an image to detect faces using OpenCV's ONNX-based YuNet face detector."
)
if __name__ == "__main__":
demo.launch()