from __future__ import print_function
import numpy as np
import torch
import torch.backends.cudnn as cudnn
from retinaface.data import cfg_mnet
from retinaface.layers.functions.prior_box import PriorBox
from retinaface.loader import load_model
from retinaface.utils.box_utils import decode, decode_landm
from retinaface.utils.nms.py_cpu_nms import py_cpu_nms
cudnn.benchmark = True
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# device = torch.device('cpu')
model = load_model().to(device)
model.eval()
def detect_faces(img_raw, confidence_threshold=0.9, top_k=5000, nms_threshold=0.4, keep_top_k=750, resize=1):
img = np.float32(img_raw)
im_height, im_width = img.shape[:2]
scale = torch.Tensor([img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
img -= (104, 117, 123)
img = img.transpose(2, 0, 1)
img = torch.from_numpy(img).unsqueeze(0)
img = img.to(device)
scale = scale.to(device)
# tic = time.time()
with torch.no_grad():
def __init__(self, net='mnet', type='cuda'):
cudnn.benchmark = True
self.net = net
self.device = torch.device(type)
self.model = load_model(net).to(self.device)
self.model.eval()