Spaces:
Runtime error
Runtime error
| from __future__ import division | |
| import numpy as np | |
| import cv2 | |
| from ..utils.load_model import load_model | |
| def distance2bbox(points, distance, max_shape=None): | |
| """Decode distance prediction to bounding box. | |
| Args: | |
| points (Tensor): Shape (n, 2), [x, y]. | |
| distance (Tensor): Distance from the given point to 4 | |
| boundaries (left, top, right, bottom). | |
| max_shape (tuple): Shape of the image. | |
| Returns: | |
| Tensor: Decoded bboxes. | |
| """ | |
| x1 = points[:, 0] - distance[:, 0] | |
| y1 = points[:, 1] - distance[:, 1] | |
| x2 = points[:, 0] + distance[:, 2] | |
| y2 = points[:, 1] + distance[:, 3] | |
| if max_shape is not None: | |
| x1 = x1.clamp(min=0, max=max_shape[1]) | |
| y1 = y1.clamp(min=0, max=max_shape[0]) | |
| x2 = x2.clamp(min=0, max=max_shape[1]) | |
| y2 = y2.clamp(min=0, max=max_shape[0]) | |
| return np.stack([x1, y1, x2, y2], axis=-1) | |
| def distance2kps(points, distance, max_shape=None): | |
| """Decode distance prediction to bounding box. | |
| Args: | |
| points (Tensor): Shape (n, 2), [x, y]. | |
| distance (Tensor): Distance from the given point to 4 | |
| boundaries (left, top, right, bottom). | |
| max_shape (tuple): Shape of the image. | |
| Returns: | |
| Tensor: Decoded bboxes. | |
| """ | |
| preds = [] | |
| for i in range(0, distance.shape[1], 2): | |
| px = points[:, i%2] + distance[:, i] | |
| py = points[:, i%2+1] + distance[:, i+1] | |
| if max_shape is not None: | |
| px = px.clamp(min=0, max=max_shape[1]) | |
| py = py.clamp(min=0, max=max_shape[0]) | |
| preds.append(px) | |
| preds.append(py) | |
| return np.stack(preds, axis=-1) | |
| class InsightFaceDet: | |
| def __init__(self, model_path, device="cuda", **kwargs): | |
| kwargs["model_file"] = model_path | |
| kwargs["module_name"] = "RetinaFace" | |
| kwargs["package_name"] = "..aux_models.modules" | |
| self.model, self.model_type = load_model(model_path, device=device, **kwargs) | |
| self.device = device | |
| if self.model_type != "ori": | |
| self._init_vars() | |
| def _init_vars(self): | |
| self.center_cache = {} | |
| self.nms_thresh = 0.4 | |
| self.det_thresh = 0.5 | |
| self.input_size = (512, 512) | |
| self.input_mean = 127.5 | |
| self.input_std = 128.0 | |
| self._anchor_ratio = 1.0 | |
| self.fmc = 3 | |
| self._feat_stride_fpn = [8, 16, 32] | |
| self._num_anchors = 2 | |
| self.use_kps = True | |
| self.output_names = [ | |
| "scores1", | |
| "scores2", | |
| "scores3", | |
| "boxes1", | |
| "boxes2", | |
| "boxes3", | |
| "kps1", | |
| "kps2", | |
| "kps3", | |
| ] | |
| def _run_model(self, blob): | |
| if self.model_type == "onnx": | |
| net_outs = self.model.run(None, {"image": blob}) | |
| elif self.model_type == "tensorrt": | |
| self.model.setup({"image": blob}) | |
| self.model.infer() | |
| net_outs = [self.model.buffer[name][0] for name in self.output_names] | |
| else: | |
| raise ValueError(f"Unsupported model type: {self.model_type}") | |
| return net_outs | |
| def _forward(self, img, threshold): | |
| """ | |
| img: np.ndarray, shape (h, w, 3) | |
| """ | |
| scores_list = [] | |
| bboxes_list = [] | |
| kpss_list = [] | |
| input_size = tuple(img.shape[0:2][::-1]) | |
| blob = cv2.dnn.blobFromImage(img, 1.0/self.input_std, input_size, (self.input_mean, self.input_mean, self.input_mean), swapRB=True) | |
| # (1, 3, 512, 512) | |
| net_outs = self._run_model(blob) | |
| input_height = blob.shape[2] | |
| input_width = 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] | |
| bbox_preds = bbox_preds * stride | |
| if self.use_kps: | |
| kps_preds = net_outs[idx+fmc*2] * stride | |
| height = input_height // stride | |
| width = input_width // stride | |
| # K = height * width | |
| key = (height, width, stride) | |
| if key in self.center_cache: | |
| anchor_centers = self.center_cache[key] | |
| else: | |
| #solution-3: | |
| 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>=threshold)[0] | |
| bboxes = distance2bbox(anchor_centers, bbox_preds) | |
| pos_scores = scores[pos_inds] | |
| pos_bboxes = bboxes[pos_inds] | |
| scores_list.append(pos_scores) | |
| bboxes_list.append(pos_bboxes) | |
| if self.use_kps: | |
| kpss = distance2kps(anchor_centers, kps_preds) | |
| kpss = kpss.reshape( (kpss.shape[0], -1, 2) ) | |
| pos_kpss = kpss[pos_inds] | |
| kpss_list.append(pos_kpss) | |
| return scores_list, bboxes_list, kpss_list | |
| def detect(self, img, input_size=None, max_num=0, metric='default', det_thresh=None): | |
| input_size = self.input_size if input_size is None else input_size | |
| det_thresh = self.det_thresh if det_thresh is None else det_thresh | |
| 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_img = 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_img | |
| scores_list, bboxes_list, kpss_list = self._forward(det_img, det_thresh) | |
| scores = np.vstack(scores_list) | |
| scores_ravel = scores.ravel() | |
| order = scores_ravel.argsort()[::-1] | |
| bboxes = np.vstack(bboxes_list) / det_scale | |
| if self.use_kps: | |
| kpss = np.vstack(kpss_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, :] | |
| if self.use_kps: | |
| kpss = kpss[order,:,:] | |
| kpss = kpss[keep,:,:] | |
| else: | |
| kpss = None | |
| 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) | |
| if metric=='max': | |
| values = area | |
| else: | |
| values = area - offset_dist_squared * 2.0 # some extra weight on the centering | |
| bindex = np.argsort(values)[::-1] # some extra weight on the centering | |
| bindex = bindex[0:max_num] | |
| det = det[bindex, :] | |
| if kpss is not None: | |
| kpss = kpss[bindex, :] | |
| return det, kpss | |
| def nms(self, dets): | |
| thresh = self.nms_thresh | |
| x1 = dets[:, 0] | |
| y1 = dets[:, 1] | |
| x2 = dets[:, 2] | |
| y2 = dets[:, 3] | |
| scores = 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 <= thresh)[0] | |
| order = order[inds + 1] | |
| return keep | |
| def __call__(self, img, **kwargs): | |
| if self.model_type == "ori": | |
| det, kpss = self.model.detect(img, **kwargs) | |
| else: | |
| det, kpss = self.detect(img, **kwargs) | |
| return det, kpss | |