Spaces:
Runtime error
Runtime error
File size: 5,765 Bytes
f2dbf59 |
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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
import os
import onnx
import onnxruntime
import cv2
import numpy as np
from collections import namedtuple
Face = namedtuple('Face',
[
'bounding_box',
'landmarks',
'scores',
])
class YoloFaceOnnx:
def __init__(self, model_path, providers):
self.session = onnxruntime.InferenceSession(model_path, providers=providers)
inputs = self.session.get_inputs()
self.input_size = (inputs[0].shape[2], inputs[0].shape[3])
self.input_name = inputs[0].name
def pre_process(self, image):
img = cv2.resize(image, self.input_size)
img = (img - 127.5) / 128.0
img = np.expand_dims(img.transpose(2, 0, 1), axis = 0).astype(np.float32)
return img
def post_process(self, size, bounding_box_list, face_landmark_5_list, score_list):
sort_indices = np.argsort(-np.array(score_list))
bounding_box_list = [ bounding_box_list[index] for index in sort_indices ]
face_landmark_5_list = [face_landmark_5_list[index] for index in sort_indices]
score_list = [ score_list[index] for index in sort_indices ]
face_list = []
keep_indices = self.apply_nms(bounding_box_list, 0.4)
for index in keep_indices:
bounding_box = bounding_box_list[index]
face_landmark = face_landmark_5_list[index]
score = score_list[index],
#print(f'bounding_box >> : {bounding_box}')
face_list.append(Face(
bounding_box = self.expand_bounding_box(size, bounding_box),
landmarks = face_landmark,
scores = score,
))
return face_list
def apply_nms(self, bounding_box_list, iou_threshold):
keep_indices = []
dimension_list = np.reshape(bounding_box_list, (-1, 4))
x1 = dimension_list[:, 0]
y1 = dimension_list[:, 1]
x2 = dimension_list[:, 2]
y2 = dimension_list[:, 3]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
indices = np.arange(len(bounding_box_list))
while indices.size > 0:
index = indices[0]
remain_indices = indices[1:]
keep_indices.append(index)
xx1 = np.maximum(x1[index], x1[remain_indices])
yy1 = np.maximum(y1[index], y1[remain_indices])
xx2 = np.minimum(x2[index], x2[remain_indices])
yy2 = np.minimum(y2[index], y2[remain_indices])
width = np.maximum(0, xx2 - xx1 + 1)
height = np.maximum(0, yy2 - yy1 + 1)
iou = width * height / (areas[index] + areas[remain_indices] - width * height)
indices = indices[np.where(iou <= iou_threshold)[0] + 1]
return keep_indices
def expand_bounding_box(self, size, bbox, expansion=32):
"""
根据给定的 bbox，在四个边各扩展一定像素。
:param image: 输入的图像
:param bbox: 原始边界框，格式为 (x1, y1, x2, y2)
:param expansion: 每个边扩展的像素数
:return: 扩展后的图像
"""
x1, y1, x2, y2 = bbox
# 扩展边界框
x1_expanded = max(x1 - expansion, 0) # 保证不越界
y1_expanded = max(y1 - expansion, 0)
x2_expanded = min(x2 + expansion, size[1]) # 保证不越界
y2_expanded = min(y2 + expansion, size[0])
return [x1_expanded, y1_expanded, x2_expanded, y2_expanded]
def detect(self, image, conf):
img = self.pre_process(image)
ratio_height = image.shape[0] / img.shape[2]
ratio_width = image.shape[1] / img.shape[3]
outputs = self.session.run(None, {self.input_name: img})
outputs = np.squeeze(outputs).T
bounding_box_raw, score_raw, face_landmark_5_raw = np.split(outputs, [ 4, 5 ], axis = 1)
bounding_box_list = []
face_landmark_5_list = []
score_list = []
keep_indices = np.where(score_raw > conf)[0]
if keep_indices.any():
bounding_box_raw, face_landmark_5_raw, score_raw = bounding_box_raw[keep_indices], face_landmark_5_raw[keep_indices], score_raw[keep_indices]
for bounding_box in bounding_box_raw:
bounding_box_list.append(np.array(
[
(bounding_box[0] - bounding_box[2] / 2) * ratio_width,
(bounding_box[1] - bounding_box[3] / 2) * ratio_height,
(bounding_box[0] + bounding_box[2] / 2) * ratio_width,
(bounding_box[1] + bounding_box[3] / 2) * ratio_height
]))
face_landmark_5_raw[:, 0::3] = (face_landmark_5_raw[:, 0::3]) * ratio_width
face_landmark_5_raw[:, 1::3] = (face_landmark_5_raw[:, 1::3]) * ratio_height
for face_landmark_5 in face_landmark_5_raw:
face_landmark_5_list.append(np.array(face_landmark_5.reshape(-1, 3)[:, :2]))
score_list = score_raw.ravel().tolist()
return self.post_process(image.shape, bounding_box_list, face_landmark_5_list, score_list)
if __name__ == "__main__":
model_path = '/Users/wadahana/workspace/AI/sd/ComfyUI/models/facefusion/yoloface_8n.onnx'
providers=['CPUExecutionProvider']
detector = YoloFaceOnnx(model_path=model_path, providers=providers)
image = cv2.imread('/Users/wadahana/Desktop/test4.jpg')
face_list = detector.detect(image=image, conf=0.7)
face = face_list[0]
pt1 = (int(face.bounding_box[0]), int(face.bounding_box[1])) # 左上角 (x1, y1)
pt2 = (int(face.bounding_box[2]), int(face.bounding_box[3])) # 右下角 (x2, y2)
cv2.rectangle(image, pt1, pt2, (255, 0, 0), 1)
print(f'face_list: {face_list}')
cv2.imwrite('/Users/wadahana/Desktop/output.jpg', image)
|