import numpy as np import onnx, onnx2torch, cv2 import torch from insightface.utils import face_align class ArcFaceRecognizer: def __init__(self, model_file=None, device='cpu', dtype=torch.float32): assert model_file is not None self.model_file = model_file self.device = device self.dtype = dtype self.model = onnx2torch.convert(onnx.load(model_file)).to(device=device, dtype=dtype) for param in self.model.parameters(): param.requires_grad = False self.model.eval() self.input_mean = 127.5 self.input_std = 127.5 self.input_size = (112, 112) self.input_shape = ['None', 3, 112, 112] def get(self, img, face): aimg = face_align.norm_crop(img, landmark=face.kps, image_size=self.input_size[0]) face.embedding = self.get_feat(aimg).flatten() return face.embedding def compute_sim(self, feat1, feat2): from numpy.linalg import norm feat1 = feat1.ravel() feat2 = feat2.ravel() sim = np.dot(feat1, feat2) / (norm(feat1) * norm(feat2)) return sim def get_feat(self, imgs): if not isinstance(imgs, list): imgs = [imgs] input_size = self.input_size blob = cv2.dnn.blobFromImages(imgs, 1.0 / self.input_std, input_size, (self.input_mean, self.input_mean, self.input_mean), swapRB=True) blob_torch = torch.tensor(blob).to(device=self.device, dtype=self.dtype) net_out = self.model(blob_torch) return net_out[0].float().cpu() 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 FaceDetector: def __init__(self, model_file=None, dtype=torch.float32, device='cuda'): self.model_file = model_file self.taskname = 'detection' self.center_cache = {} self.nms_thresh = 0.4 self.det_thresh = 0.5 self.device = device self.dtype = dtype self.model = onnx2torch.convert(onnx.load(model_file)).to(device=device, dtype=dtype) for param in self.model.parameters(): param.requires_grad = False self.model.eval() input_shape = (320, 320) self.input_size = input_shape self.input_shape = input_shape self.input_mean = 127.5 self.input_std = 128.0 self._anchor_ratio = 1.0 self._num_anchors = 1 self.fmc = 3 self._feat_stride_fpn = [8, 16, 32] self._num_anchors = 2 self.use_kps = True self.det_thresh = 0.5 self.nms_thresh = 0.4 def forward(self, img, threshold): 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) blob_torch = torch.tensor(blob).to(device=self.device, dtype=self.dtype) net_outs_torch = self.model(blob_torch) # print(list(map(lambda x: x.shape, net_outs_torch))) net_outs = list(map(lambda x: x.float().cpu().numpy(), net_outs_torch)) 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-1, c style: # anchor_centers = np.zeros( (height, width, 2), dtype=np.float32 ) # for i in range(height): # anchor_centers[i, :, 1] = i # for i in range(width): # anchor_centers[:, i, 0] = i # solution-2: # ax = np.arange(width, dtype=np.float32) # ay = np.arange(height, dtype=np.float32) # xv, yv = np.meshgrid(np.arange(width), np.arange(height)) # anchor_centers = np.stack([xv, yv], axis=-1).astype(np.float32) # solution-3: anchor_centers = np.stack(np.mgrid[:height, :width][::-1], axis=-1).astype(np.float32) # print(anchor_centers.shape) 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 = 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'): assert input_size is not None or self.input_size is not None input_size = self.input_size if input_size is None else input_size 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, self.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