Spaces:
Running
Running
| from third_party.PIPNet.FaceBoxesV2.detector import Detector | |
| import cv2, os | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from third_party.PIPNet.FaceBoxesV2.utils.config import cfg | |
| from third_party.PIPNet.FaceBoxesV2.utils.prior_box import PriorBox | |
| from third_party.PIPNet.FaceBoxesV2.utils.nms_wrapper import nms | |
| from third_party.PIPNet.FaceBoxesV2.utils.faceboxes import FaceBoxesV2 | |
| from third_party.PIPNet.FaceBoxesV2.utils.box_utils import decode | |
| import time | |
| class FaceBoxesDetector(Detector): | |
| def __init__(self, model_arch, model_weights, use_gpu, device): | |
| super().__init__(model_arch, model_weights) | |
| self.name = "FaceBoxesDetector" | |
| self.net = FaceBoxesV2( | |
| phase="test", size=None, num_classes=2 | |
| ) # initialize detector | |
| self.use_gpu = use_gpu | |
| self.device = device | |
| state_dict = torch.load(self.model_weights, map_location=self.device) | |
| # create new OrderedDict that does not contain `module.` | |
| from collections import OrderedDict | |
| new_state_dict = OrderedDict() | |
| for k, v in state_dict.items(): | |
| name = k[7:] # remove `module.` | |
| new_state_dict[name] = v | |
| # load params | |
| self.net.load_state_dict(new_state_dict) | |
| self.net = self.net.to(self.device) | |
| self.net.eval() | |
| def detect(self, image, thresh=0.6, im_scale=None): | |
| # auto resize for large images | |
| if im_scale is None: | |
| height, width, _ = image.shape | |
| if min(height, width) > 600: | |
| im_scale = 600.0 / min(height, width) | |
| else: | |
| im_scale = 1 | |
| image_scale = cv2.resize( | |
| image, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR | |
| ) | |
| scale = torch.Tensor( | |
| [ | |
| image_scale.shape[1], | |
| image_scale.shape[0], | |
| image_scale.shape[1], | |
| image_scale.shape[0], | |
| ] | |
| ) | |
| image_scale = ( | |
| torch.from_numpy(image_scale.transpose(2, 0, 1)).to(self.device).int() | |
| ) | |
| mean_tmp = torch.IntTensor([104, 117, 123]).to(self.device) | |
| mean_tmp = mean_tmp.unsqueeze(1).unsqueeze(2) | |
| image_scale -= mean_tmp | |
| image_scale = image_scale.float().unsqueeze(0) | |
| scale = scale.to(self.device) | |
| with torch.no_grad(): | |
| out = self.net(image_scale) | |
| # priorbox = PriorBox(cfg, out[2], (image_scale.size()[2], image_scale.size()[3]), phase='test') | |
| priorbox = PriorBox( | |
| cfg, image_size=(image_scale.size()[2], image_scale.size()[3]) | |
| ) | |
| priors = priorbox.forward() | |
| priors = priors.to(self.device) | |
| loc, conf = out | |
| prior_data = priors.data | |
| boxes = decode(loc.data.squeeze(0), prior_data, cfg["variance"]) | |
| boxes = boxes * scale | |
| boxes = boxes.cpu().numpy() | |
| scores = conf.data.cpu().numpy()[:, 1] | |
| # ignore low scores | |
| inds = np.where(scores > thresh)[0] | |
| boxes = boxes[inds] | |
| scores = scores[inds] | |
| # keep top-K before NMS | |
| order = scores.argsort()[::-1][:5000] | |
| boxes = boxes[order] | |
| scores = scores[order] | |
| # do NMS | |
| dets = np.hstack((boxes, scores[:, np.newaxis])).astype( | |
| np.float32, copy=False | |
| ) | |
| keep = nms(dets, 0.3) | |
| dets = dets[keep, :] | |
| dets = dets[:750, :] | |
| detections_scale = [] | |
| for i in range(dets.shape[0]): | |
| xmin = int(dets[i][0]) | |
| ymin = int(dets[i][1]) | |
| xmax = int(dets[i][2]) | |
| ymax = int(dets[i][3]) | |
| score = dets[i][4] | |
| width = xmax - xmin | |
| height = ymax - ymin | |
| detections_scale.append(["face", score, xmin, ymin, width, height]) | |
| # adapt bboxes to the original image size | |
| if len(detections_scale) > 0: | |
| detections_scale = [ | |
| [ | |
| det[0], | |
| det[1], | |
| int(det[2] / im_scale), | |
| int(det[3] / im_scale), | |
| int(det[4] / im_scale), | |
| int(det[5] / im_scale), | |
| ] | |
| for det in detections_scale | |
| ] | |
| return detections_scale, im_scale | |