Spaces:
Sleeping
Sleeping
| """Standalone SCRFD face detector (ONNX Runtime, CPU). | |
| A self-contained re-implementation of InsightFace's SCRFD post-processing so the | |
| app can load a raw ``.onnx`` file directly — no ``insightface`` package, no model | |
| pack directory. Handles the 6/9/10/15-output SCRFD variants; ``det_500m.onnx`` and | |
| ``det_2.5g.onnx`` are both 9-output (3 strides x {score, bbox, kps}) with 5-point | |
| landmarks. | |
| Reference: https://github.com/deepinsight/insightface (scrfd.py). | |
| """ | |
| import cv2 | |
| import numpy as np | |
| import onnxruntime | |
| def distance2bbox(points, distance): | |
| """Decode (left, top, right, bottom) distances from an anchor center to a box.""" | |
| x1 = points[:, 0] - distance[:, 0] | |
| y1 = points[:, 1] - distance[:, 1] | |
| x2 = points[:, 0] + distance[:, 2] | |
| y2 = points[:, 1] + distance[:, 3] | |
| return np.stack([x1, y1, x2, y2], axis=-1) | |
| def distance2kps(points, distance): | |
| """Decode per-keypoint (dx, dy) distances from an anchor center to landmarks.""" | |
| preds = [] | |
| for i in range(0, distance.shape[1], 2): | |
| px = points[:, i % 2] + distance[:, i] | |
| py = points[:, i % 2 + 1] + distance[:, i + 1] | |
| preds.append(px) | |
| preds.append(py) | |
| return np.stack(preds, axis=-1) | |
| class SCRFD: | |
| """SCRFD ONNX face detector. | |
| Parameters | |
| ---------- | |
| model_file : str | |
| Path to the ``.onnx`` file. | |
| providers : list[str] | None | |
| ONNX Runtime execution providers. Defaults to CPU. | |
| """ | |
| def __init__(self, model_file, providers=None): | |
| self.model_file = model_file | |
| providers = providers or ["CPUExecutionProvider"] | |
| self.session = onnxruntime.InferenceSession(model_file, providers=providers) | |
| self.center_cache = {} | |
| self.nms_thresh = 0.4 | |
| self.input_mean = 127.5 | |
| self.input_std = 128.0 | |
| self._init_vars() | |
| def _init_vars(self): | |
| inp = self.session.get_inputs()[0] | |
| self.input_name = inp.name | |
| # Static input size if the model declares one (e.g. [1, 3, 640, 640]), | |
| # else fall back to 640x640 at detect() time. | |
| shape = inp.shape | |
| if isinstance(shape[2], int) and isinstance(shape[3], int): | |
| self.input_size = (shape[3], shape[2]) # (w, h) | |
| else: | |
| self.input_size = (640, 640) | |
| outputs = self.session.get_outputs() | |
| self.output_names = [o.name for o in outputs] | |
| self.use_kps = False | |
| self._num_anchors = 1 | |
| n = len(outputs) | |
| if n == 6: | |
| self.fmc, self._feat_stride_fpn, self._num_anchors = 3, [8, 16, 32], 2 | |
| elif n == 9: | |
| self.fmc, self._feat_stride_fpn, self._num_anchors = 3, [8, 16, 32], 2 | |
| self.use_kps = True | |
| elif n == 10: | |
| self.fmc, self._feat_stride_fpn, self._num_anchors = 5, [8, 16, 32, 64, 128], 1 | |
| elif n == 15: | |
| self.fmc, self._feat_stride_fpn, self._num_anchors = 5, [8, 16, 32, 64, 128], 1 | |
| self.use_kps = True | |
| else: | |
| raise ValueError(f"Unexpected SCRFD output count: {n}") | |
| def forward(self, img, thresh): | |
| scores_list, bboxes_list, kpss_list = [], [], [] | |
| blob = cv2.dnn.blobFromImage( | |
| img, | |
| 1.0 / self.input_std, | |
| (img.shape[1], img.shape[0]), | |
| (self.input_mean, self.input_mean, self.input_mean), | |
| swapRB=True, | |
| ) | |
| net_outs = self.session.run(self.output_names, {self.input_name: blob}) | |
| input_height, input_width = blob.shape[2], blob.shape[3] | |
| fmc = self.fmc | |
| for idx, stride in enumerate(self._feat_stride_fpn): | |
| scores = net_outs[idx] | |
| bbox_preds = net_outs[idx + fmc] * stride | |
| if self.use_kps: | |
| kps_preds = net_outs[idx + fmc * 2] * stride | |
| height, width = input_height // stride, input_width // stride | |
| key = (height, width, stride) | |
| if key in self.center_cache: | |
| anchor_centers = self.center_cache[key] | |
| else: | |
| anchor_centers = np.stack( | |
| np.mgrid[:height, :width][::-1], axis=-1 | |
| ).astype(np.float32) | |
| anchor_centers = (anchor_centers * stride).reshape((-1, 2)) | |
| if self._num_anchors > 1: | |
| anchor_centers = np.stack( | |
| [anchor_centers] * self._num_anchors, axis=1 | |
| ).reshape((-1, 2)) | |
| if len(self.center_cache) < 100: | |
| self.center_cache[key] = anchor_centers | |
| pos_inds = np.where(scores >= thresh)[0] | |
| bboxes = distance2bbox(anchor_centers, bbox_preds) | |
| scores_list.append(scores[pos_inds]) | |
| bboxes_list.append(bboxes[pos_inds]) | |
| if self.use_kps: | |
| kpss = distance2kps(anchor_centers, kps_preds) | |
| kpss = kpss.reshape((kpss.shape[0], -1, 2)) | |
| kpss_list.append(kpss[pos_inds]) | |
| return scores_list, bboxes_list, kpss_list | |
| def nms(self, dets): | |
| """Standard IoU NMS on ``[x1, y1, x2, y2, score]`` rows (score-sorted).""" | |
| x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4] | |
| areas = (x2 - x1 + 1) * (y2 - y1 + 1) | |
| order = scores.argsort()[::-1] | |
| keep = [] | |
| while order.size > 0: | |
| i = order[0] | |
| keep.append(i) | |
| xx1 = np.maximum(x1[i], x1[order[1:]]) | |
| yy1 = np.maximum(y1[i], y1[order[1:]]) | |
| xx2 = np.minimum(x2[i], x2[order[1:]]) | |
| yy2 = np.minimum(y2[i], y2[order[1:]]) | |
| w = np.maximum(0.0, xx2 - xx1 + 1) | |
| h = np.maximum(0.0, yy2 - yy1 + 1) | |
| inter = w * h | |
| ovr = inter / (areas[i] + areas[order[1:]] - inter) | |
| inds = np.where(ovr <= self.nms_thresh)[0] | |
| order = order[inds + 1] | |
| return keep | |
| def detect(self, img, thresh=0.5, input_size=None, max_num=0, metric="default"): | |
| """Detect faces in a BGR ``uint8`` image. | |
| Returns ``(dets, kpss)`` where ``dets`` is ``[N, 5]`` (x1,y1,x2,y2,score) | |
| in original-image pixels and ``kpss`` is ``[N, 5, 2]`` landmarks (or None). | |
| """ | |
| input_size = self.input_size if input_size is None else input_size | |
| # Letterbox: preserve aspect ratio, pad to the model's square input. | |
| im_ratio = float(img.shape[0]) / img.shape[1] | |
| model_ratio = float(input_size[1]) / input_size[0] | |
| if im_ratio > model_ratio: | |
| new_height = input_size[1] | |
| new_width = int(new_height / im_ratio) | |
| else: | |
| new_width = input_size[0] | |
| new_height = int(new_width * im_ratio) | |
| det_scale = float(new_height) / img.shape[0] | |
| resized = cv2.resize(img, (new_width, new_height)) | |
| det_img = np.zeros((input_size[1], input_size[0], 3), dtype=np.uint8) | |
| det_img[:new_height, :new_width, :] = resized | |
| scores_list, bboxes_list, kpss_list = self.forward(det_img, thresh) | |
| scores = np.vstack(scores_list) | |
| order = scores.ravel().argsort()[::-1] | |
| bboxes = np.vstack(bboxes_list) / det_scale | |
| pre_det = np.hstack((bboxes, scores)).astype(np.float32, copy=False) | |
| pre_det = pre_det[order, :] | |
| keep = self.nms(pre_det) | |
| det = pre_det[keep, :] | |
| kpss = None | |
| if self.use_kps: | |
| kpss = np.vstack(kpss_list) / det_scale | |
| kpss = kpss[order, :, :][keep, :, :] | |
| if max_num > 0 and det.shape[0] > max_num: | |
| area = (det[:, 2] - det[:, 0]) * (det[:, 3] - det[:, 1]) | |
| img_center = img.shape[0] // 2, img.shape[1] // 2 | |
| offsets = np.vstack([ | |
| (det[:, 0] + det[:, 2]) / 2 - img_center[1], | |
| (det[:, 1] + det[:, 3]) / 2 - img_center[0], | |
| ]) | |
| offset_dist_squared = np.sum(np.power(offsets, 2.0), 0) | |
| values = area if metric == "max" else area - offset_dist_squared * 2.0 | |
| bindex = np.argsort(values)[::-1][:max_num] | |
| det = det[bindex, :] | |
| if kpss is not None: | |
| kpss = kpss[bindex, :, :] | |
| return det, kpss | |