Spaces:
Runtime error
Runtime error
| 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) | |