def main(): args = get_args() torch.set_grad_enabled(False) cfg = None if args.network == "mobile0.25": cfg = cfg_mnet elif args.network == "resnet50": cfg = cfg_re50 # net and model net = RetinaFace(cfg=cfg, phase="test") net = load_model(net, args.trained_model, args.cpu) net.eval() print("Finished loading model!") print(net) device = torch.device("cpu" if args.cpu else "cuda") net = net.to(device) # ------------------------ export ----------------------------- output_onnx = "FaceDetector.onnx" print("==> Exporting model to ONNX format at '{}'".format(output_onnx)) input_names = ["input0"] output_names = ["output0"] inputs = torch.randn(1, 3, args.long_side, args.long_side).to(device) torch.onnx._export(net, inputs, output_onnx, export_params=True, verbose=False, input_names=input_names, output_names=output_names)
def __init__(self, device, pretrained_model): self.device = device self.cfg = { 'name': 'Resnet50', 'min_sizes': [[16, 32], [64, 128], [256, 512]], 'steps': [8, 16, 32], 'variance': [0.1, 0.2], 'clip': False, 'loc_weight': 2.0, 'gpu_train': True, 'batch_size': 24, 'ngpu': 4, 'epoch': 100, 'decay1': 70, 'decay2': 90, 'image_size': 840, 'pretrain': True, 'return_layers': { 'layer2': 1, 'layer3': 2, 'layer4': 3 }, 'in_channel': 256, 'out_channel': 256 } self.net = RetinaFace(cfg=self.cfg, phase='test') self.net = load_model(self.net, pretrained_model, device) self.net.eval() self.net = self.net.to(device)
def validate(model_path, network='mobile0.25'): cfg = None if network == "mobile0.25": cfg = cfg_mnet elif network == "resnet50": cfg = cfg_re50 else: raise ValueError(network) net = RetinaFace(cfg=cfg, phase='test') net = load_model(net, model_path, args.cpu) #net.eval() print('Finished loading model!') #print(net) cudnn.benchmark = True device = torch.device("cpu" if args.cpu else "cuda") net = net.to(device) net.eval() #net.phase = 'eval' with torch.no_grad(): preds = predict(net, cfg) #net.phase = 'train' #net.train() del net aps = evaluation(preds, './widerface_evaluate/ground_truth/') avg = np.mean(aps) return [avg] + aps
def init_detector(cfg: Dict[str, any], weights: str, device: torch.device) -> torch.nn.Module: cfg['pretrain'] = False net = RetinaFace(cfg=cfg, phase='test') net = load_model(net, weights, device) net.eval() return net
def __init__(self, network, confidence=0.02, top_k=5000, nms_thresh=0.4, keep_top_k=750, vis_thresh=0.6): torch.set_grad_enabled(False) self.confidence = confidence self.top_k = top_k self.nms_thresh = nms_thresh self.keep_top_k = keep_top_k self.vis_thresh = vis_thresh self.device = "cuda" if torch.cuda.is_available() else "cpu" if network == "resnet": self.cfg = cfg_re50 model_path = "weights/Resnet50_Final.pth" else: self.cfg = cfg_mnet model_path = "weights/mobilenet0.25_Final.pth" self.net = RetinaFace(cfg=self.cfg, phase='test') self.net = load_model(self.net, model_path, True if self.device == "cpu" else False) self.net.eval() print('Finished loading model!') print(self.net) cudnn.benchmark = True self.device = torch.device(self.device) self.net = self.net.to(self.device)
def __init__( self, enable_cuda=settings.CUDA_ENABLED, face_rect_expand_factor=FACE_RECT_EXPAND_FACTOR, trained_model=settings.FACE_DETECTION_MODEL, network=settings.FACE_DETECTION_NETWORK, ): """ Initializes the RetinaFace in PyTorch Arguments: enable_cuda: boolean indicating whether CUDA must be used for the extraction of the features face_rect_expand_factor: Expansion factor for the detection face rectangle trained_model: Path to a pretrained model file with weights network: Name of the network used for the detection. The options are 'mobile0.25' or 'resnet50'. """ torch.set_grad_enabled(False) cudnn.benchmark = True self.is_cuda_enable = enable_cuda self.face_rect_expand_factor = face_rect_expand_factor self.trained_model = trained_model self.cfg = None if network == 'mobile0.25': self.cfg = cfg_mnet elif network == 'resnet50': self.cfg = cfg_re50 assert self.cfg != None, "Network name can only be 'resnet50' or 'mobile0.25' !" self.net = RetinaFace(cfg=self.cfg, phase='test') self.net = self.load_model(self.net, self.trained_model, not self.is_cuda_enable) self.net.eval() self.device = torch.device( 'cpu' if not self.is_cuda_enable else 'cuda') self.net = self.net.to(self.device)
def __init__(self, on_gpu=False): self.on_gpu = on_gpu # from classifier by Sizykh Ivan self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") # parser = argparse.ArgumentParser(description='Retinaface') # device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # # parser.add_argument('-m', '--trained_model', default='./weights/Resnet50_Final.pth', # type=str, help='Trained state_dict file path to open') # parser.add_argument('--network', default='resnet50', help='Backbone network mobile0.25 or resnet50') # # parser.add_argument('--cpu', action="store_true", default=False, help='Use cpu inference') # parser.add_argument('--cpu', action="store_true", default=False, help='Use cpu inference') # parser.add_argument('--confidence_threshold', default=0.02, type=float, help='confidence_threshold') # parser.add_argument('--top_k', default=5000, type=int, help='top_k') # parser.add_argument('--nms_threshold', default=0.4, type=float, help='nms_threshold') # parser.add_argument('--keep_top_k', default=750, type=int, help='keep_top_k') # parser.add_argument('-s', '--save_image', action="store_true", default=True, help='show detection results') # parser.add_argument('--vis_thres', default=0.6, type=float, help='visualization_threshold') # # parser.add_argument('-v', '--video', default='vid.mp4', type=str) # # args = parser.parse_args() self.trained_model = './weights/Resnet50_Final.pth' self.network = 'resnet50' self.cpu = False self.confidence_threshold = 0.02 self.top_k = 5000 self.nms_threshold = 0.4 self.keep_top_k = 750 self.vis_thres = 0.6 self.resize = 1 torch.set_grad_enabled(False) cfg = None if self.network == "mobile0.25": cfg = cfg_mnet elif self.network == "resnet50": cfg = cfg_re50 # cfg = cfg_re50 # net and model detector = RetinaFace(cfg=cfg, phase='test') detector = self.load_model(model=detector, pretrained_path=self.trained_model, load_to_cpu=self.cpu) detector.eval() print('Finished loading model!') # print(detector) if self.on_gpu: cudnn.benchmark = True self.detector = detector.to(self.device) else: self.detector = detector self.cfg = cfg
def create_model(weights): model = RetinaFace(cfg=cfg_re50, phase='test') checkpoint = load(weights, map_location='cpu') ckpt = {k.replace('module.', ''): v for k, v in checkpoint.items()} model.load_state_dict(ckpt) return model
def model_cfg(trained_model, cfg, device, cpu): net = RetinaFace(cfg=cfg, phase='test') net = load_model(net, trained_model, load_to_cpu=cpu) net.eval() print('Finished loading model!') cudnn.benchmark = True net = net.to(device) return net
def __init__(self, device="cuda", confidence_threshold=0.8): self.device = device self.confidence_threshold = confidence_threshold self.cfg = cfg = cfg_re50 self.variance = cfg["variance"] cfg["pretrain"] = False self.net = RetinaFace(cfg=cfg, phase="test").to(device).eval() self.decode_param_cache = {}
def load_models(): torch.set_grad_enabled(False) net = RetinaFace(cfg=cfg_mnet, phase='test') net = load_model(net, trained_model) print('Finished loading retinaface model!') #cudnn.benchmark = True net = net.eval().to(device) resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device) print('Finished loading resnet model!') return net, resnet
def config_retina_model(): torch.set_grad_enabled(False) cfg = cfg_re50 ### Retinaface for face cropping retinaface = RetinaFace(cfg=cfg, phase='test') retinaface = load_model(retinaface, './weights/Resnet50_Final.pth', True) retinaface.eval() #print('Finished loading model!') return retinaface, cfg
def __init__(self): torch.set_grad_enabled(False) cudnn.benchmark = True self.opt=get_config() if self.opt.network == "mobile0.25": self.cfg = cfg_mnet elif self.opt.network == "resnet50": self.cfg = cfg_re50 # net and model self.net = RetinaFace(cfg=self.cfg, phase = 'test') self.net = self.load_model(self.net, self.opt.trained_model, self.opt.cpu) self.net.eval() self.net = self.net.to(self.opt.device)
def __init__(self, on_gpu=False, confidence_threshold=0.02, top_k=5000, nms_threshold=0.4, keep_top_k=750, vis_thres=0.6, network='resnet50'): self.on_gpu = on_gpu # from classifier by Sizykh Ivan self.device = torch.device( "cuda:0" if torch.cuda.is_available() else "cpu") self.network = network self.cpu = False self.confidence_threshold = confidence_threshold self.top_k = top_k self.nms_threshold = nms_threshold self.keep_top_k = keep_top_k self.vis_thres = vis_thres if network == 'resnet50': self.trained_model = './weights/Resnet50_Final.pth' else: self.trained_model = './weights/mobilenet0.25_Final.pth' self.resize = 1 torch.set_grad_enabled(False) cfg = None if self.network == "mobile0.25": cfg = cfg_mnet elif self.network == "resnet50": cfg = cfg_re50 # cfg = cfg_re50 # net and model detector = RetinaFace(cfg=cfg, phase='test') detector = self.load_model(model=detector, pretrained_path=self.trained_model, load_to_cpu=self.cpu) detector.eval() print('Finished loading model!') # print(detector) if self.on_gpu: cudnn.benchmark = True self.detector = detector.to(self.device) else: self.detector = detector self.cfg = cfg
def Model_Params(self, model_type="mobilenet", model_path="weights_trained/mobilenet0.25_Final.pth", use_gpu=True): if (model_type == "mobilenet"): self.system_dict["local"]["network"] = "mobile0.25" elif (model_type == "resnet"): self.system_dict["local"]["network"] = "resnet50" else: print("Model type not found") print("Available models - mobilenet, resnet") self.system_dict["local"]["cpu"] = not use_gpu self.system_dict["local"]["trained_model"] = model_path torch.set_grad_enabled(False) self.system_dict["local"]["cfg"] = None if self.system_dict["local"]["network"] == "mobile0.25": self.system_dict["local"]["cfg"] = cfg_mnet elif self.system_dict["local"]["network"] == "resnet50": self.system_dict["local"]["cfg"] = cfg_re50 # net and model self.system_dict["local"]["net"] = RetinaFace( cfg=self.system_dict["local"]["cfg"], phase='test') self.system_dict["local"]["net"] = self.load_model( self.system_dict["local"]["net"], self.system_dict["local"]["trained_model"], self.system_dict["local"]["cpu"]) self.system_dict["local"]["net"].eval() print('Finished loading model!') cudnn.benchmark = True self.system_dict["local"]["device"] = torch.device( "cpu" if self.system_dict["local"]["cpu"] else "cuda") self.system_dict["local"]["net"] = self.system_dict["local"]["net"].to( self.system_dict["local"]["device"])
def load_net(): torch.set_grad_enabled(False) cfg = None if args.network == "mobile0.25": cfg = cfg_mnet elif args.network == "resnet50": cfg = cfg_re50 # net and model net = RetinaFace(cfg=cfg, phase='test') net = load_model(net, args.trained_model, args.cpu) net.eval() print('Finished loading model!') cudnn.benchmark = True device = torch.device("cpu" if args.cpu else "cuda") net = net.to(device) return net, device, cfg
def _initialize_weight(self): self.cfg = None if self.network == "mobile0.25": self.cfg = cfg_mnet elif self.network == "resnet50": self.cfg = cfg_re50 self.net = RetinaFace(cfg=self.cfg, phase='test') self.net = self._load_model(self.net, self.weight_path, self.use_cpu) self.net.eval() print('Finished loading model!') print(self.net) cudnn.benchmark = True self.device = torch.device("cpu" if self.use_cpu else "cuda") print("self. device : ", self.device) self.net = self.net.to(self.device)
def __init__(self,model_type="mobile0.25",model_path="./weights/mobilenet0.25_Final.pth", backbone_location="./weights/mobilenetV1X0.25_pretrain.tar",use_cpu=True,loading=True): self.cfg = None self.use_cpu = use_cpu self.model_path = model_path if model_type == "mobile0.25": self.cfg = cfg_mnet elif model_type == "resnet50": self.cfg = cfg_re50 self.device = torch.device("cpu" if use_cpu else "cuda") self.net = RetinaFace(cfg=self.cfg,phase="test",backbone_location=backbone_location) if loading: print('No model path exist!') if not os.path.exists(model_path) else None self.loading() self._priors = None self.im_width = 0 self.im_height = 0 self.im_nch = 0
def initFaceDetector(): global args, net, device, cfg init_args() torch.set_grad_enabled(False) cfg = None if args.network == "mobile0.25": cfg = cfg_mnet elif args.network == "resnet50": cfg = cfg_re50 net = RetinaFace(cfg=cfg, phase='test') net, forest = load_model(net, args.trained_model, args.cpu, args.forest_path) net.eval() cudnn.benchmark = True device = torch.device("cpu" if args.cpu else "cuda") net = net.to(device)
def __init__(self, trained_model='weights/Resnet50_Final.pth', network='resnet50', cpu=False, confidence_threshold=0.02, top_k=5000, nms_threshold=0.4, keep_top_k=750, show_image=False, vis_thres=0.6 ): self.cpu = cpu self.confidence_threshold = confidence_threshold self.top_k = top_k self.nms_threshold = nms_threshold self.keep_top_k = keep_top_k self.show_image = show_image self.vis_thres = vis_thres torch.set_grad_enabled(False) self.cfg = None if network == "mobile0.25": self.cfg = cfg_mnet elif network == "resnet50": self.cfg = cfg_re50 self.cfg['pretrain'] = False # net and model self.net = RetinaFace(cfg=self.cfg, phase = 'test') self.net = load_model(self.net, os.path.join(os.path.dirname(inspect.getabsfile(RetinaFace)), '../' + trained_model), cpu) self.net.eval() print('Finished loading model!') print(self.net) cudnn.benchmark = True self.device = torch.device("cpu" if cpu else "cuda") self.net = self.net.to(self.device) self.resize = 1
def __init__(self): # TODO: add initialization logic torch.set_grad_enabled(False) self.cfg = None if args.network == "mobile0.25": self.cfg = cfg_mnet elif args.network == "resnet50": self.cfg = cfg_re50 elif args.network == "resnet18": self.cfg = cfg_re18 elif args.network == "resnet34": self.cfg = cfg_re34 # net and model self.net = RetinaFace(cfg=self.cfg, phase='test') # self.net = load_model(self.net, args.trained_model, args.cpu) self.net.eval() print('Finished loading model!') print(self.net) cudnn.benchmark = True self.device = torch.device("cpu" if args.cpu else "cuda") self.net = self.net.to(self.device) self.resize = 1
def __init__(self, threshold=0.5, network="mobile0.25", device=torch.device("cuda")): torch.set_grad_enabled(False) cfg = None if network == "mobile0.25": cfg = cfg_mnet elif network == "resnet50": cfg = cfg_re50 # net and model net = RetinaFace(cfg=cfg, phase='test') net = load_model(net, "./weights/mobilenet0.25_Final.pth", False) net.eval() # print('Finished loading model!') # print(net) cudnn.benchmark = True self.device = device net = net.to(self.device) self.net = net torch.set_grad_enabled(False) self._t = {'forward_pass': Timer(), 'misc': Timer()} self.cfg = cfg self.threshold = threshold
def __init__(self, trained_mode='./weights/mobilenet0.25_Final.pth', network='mobile0.25', cpu=True, confidence_threshold=0.02, top_k=5000, nms_threshold=0.4, keep_top_k=750, vis_thres=0.6): self.trained_model = trained_mode self.network = network self.network = network self.cpu = cpu self.confidence_threshold = confidence_threshold self.top_k = top_k self.nms_threshold = nms_threshold self.keep_top_k = keep_top_k self.vis_thres = vis_thres torch.set_grad_enabled(False) self.cfg = None if self.network == "mobile0.25": setattr(self, 'cfg', cfg_mnet) elif self.network == "resnet50": setattr(self, 'cfg', cfg_re50) else: raise (Exception("Invalid NetWork")) # build net and load model self.net = RetinaFace(self.cfg, phase='test') self.net = load_model(self.net, self.trained_model, self.cpu) self.net = self.net.eval() self.device = torch.device("cpu" if self.cpu else "cuda") self.net = self.net.to(self.device) self.resize = 1
def __init__(self, model_path, long_side, network): torch.set_grad_enabled(False) if network == 'mobilenet': self.cfg = cfg_mnet net = RetinaFace(cfg=self.cfg, phase='test') elif network == 'slim': self.cfg = cfg_slim net = Slim(cfg=self.cfg, phase='test') elif network == 'RFB': self.cfg = cfg_rfb net = RFB(cfg=self.cfg, phase='test') else: print("not supported network!!") exit(0) self.net = load_model(net, model_path, True) self.net.eval() print("Finished loading model!") cudnn.benchmark = True self.device = torch.device("cpu") self.net = self.net.to(self.device) self.long_side = long_side
else: pretrained_dict = remove_prefix(pretrained_dict, 'module.') check_keys(model, pretrained_dict) model.load_state_dict(pretrained_dict, strict=False) return model if __name__ == '__main__': torch.set_grad_enabled(False) cfg = None if args.network == "mobile0.25": cfg = cfg_mnet elif args.network == "resnet50": cfg = cfg_re50 # net and model net = RetinaFace(cfg=cfg, phase='test') net = load_model(net, args.trained_model, args.cpu) net.eval() print('Finished loading model!') print(net) cudnn.benchmark = True device = torch.device("cpu" if args.cpu else "cuda") net = net.to(device) resize = 1 for i in range(1): image_path = "nano.jpg" img_raw = cv2.imread(image_path, cv2.IMREAD_COLOR) print("The original image shape is ", img_raw.shape) img = np.float32(img_raw)
#!/usr/bin/python # -*- coding:utf-8 -*- """ @author:fangpf @time: 2020/09/28 """ import torch from torch.backends import cudnn from data import cfg_re50 from models.retinaface import RetinaFace from utils.net_utils import load_model, image_process, process_face_data cfg = cfg_re50 retina_trained_model = './weights/Resnet50_Final.pth' retina_net = RetinaFace(cfg=cfg, phase='test') retina_net = load_model(retina_net, retina_trained_model, False) retina_net = retina_net.cuda(0) retina_net.eval() cudnn.benchmark = True device = torch.device("cuda") def face_detection(im): resize = 1 im, im_width, im_height, scale = image_process(im, device) loc, conf, landms = retina_net(im) result_data = process_face_data(cfg, im, im_height, im_width, loc, scale, conf, landms, resize) return result_data
def process_video_files( network: str, trained_model: str, decode_gpu: bool, is_fp16: bool, file_paths: list, num_gpu: Optional[int], gpu_id: int, output_path: Path, is_save_boxes: bool, is_save_crops: bool, num_frames: int, resize_coeff: Optional[Tuple], confidence_threshold: float, num_workers: int, nms_threshold: float, batch_size: int, resize_scale: float, min_size: int, keep_top_k: int, ) -> None: torch.set_grad_enabled(False) if network == "mobile0.25": cfg = cfg_mnet_test elif network == "resnet50": cfg = cfg_re50_test else: raise NotImplementedError( f"Only mobile0.25 and resnet50 are suppoted, but we got {network}") if min_size < 0: raise ValueError( f"Min size should be positive, but we got {min_size}.") # net and model net = RetinaFace(cfg=cfg, phase="test") net = load_model(net, trained_model, load_to_cpu=False) net.eval() if is_fp16: net = net.half() device = torch.device("cuda") net.to(device) print("Finished loading model!") cudnn.benchmark = True transform = albu.Compose([ albu.Normalize( p=1, mean=(104, 117, 123), std=(1.0, 1.0, 1.0), max_pixel_value=1) ], p=1) if num_gpu is not None: start, end = split_array(len(file_paths), num_gpu, gpu_id) file_paths = file_paths[start:end] with torch.no_grad(): func = partial(get_frames, num_frames=num_frames, resize_coeff=resize_coeff, transform=transform, decode_gpu=decode_gpu) with torch.no_grad(): with concurrent.futures.ProcessPoolExecutor( num_workers) as executor: for result in tqdm(executor.map(func, file_paths), total=len(file_paths), leave=False, desc="Loading data files"): if len(result) != 0: result["is_fp16"] = is_fp16 result["device"] = device result["batch_size"] = batch_size result["cfg"] = cfg result["nms_threshold"] = nms_threshold result["confidence_threshold"] = confidence_threshold result["is_save_crops"] = is_save_crops result["is_save_boxes"] = is_save_boxes result["output_path"] = output_path result["net"] = net result["min_size"] = min_size result["resize_scale"] = resize_scale result["keep_top_k"] = keep_top_k process_frames(**result)
default=0.1, type=float, help='Gamma update for SGD') parser.add_argument('--save_folder', default='./weights/', help='Location to save checkpoint models') args = parser.parse_args() if not os.path.exists(args.save_folder): os.mkdir(args.save_folder) cfg = None net = None if args.network == "mobile0.25": cfg = cfg_mnet net = RetinaFace(cfg=cfg) elif args.network == "slim": cfg = cfg_slim net = Slim(cfg=cfg) elif args.network == "RFB": cfg = cfg_rfb net = RFB(cfg=cfg) elif args.network == "efficientdet": cfg = cfg_efficientdet net = EfficientDet(cfg=cfg) else: print("Don't support network!") exit(0) print("Printing net...") #print(net)
import torch import torch.backends.cudnn as cudnn import numpy as np from data import cfg from layers.functions.prior_box import PriorBox from utils.nms.py_cpu_nms import py_cpu_nms import cv2 from models.retinaface import RetinaFace from utils.box_utils import decode, decode_landm from utils.timer import Timer import numpy as np import face_recognition import pickle torch.set_grad_enabled(False) net = RetinaFace(phase="test") CONFIDENCE = 0.50 NMS_THRESHOLD = 0.30 VIZ_THRESHOLD = 0.50 resize = 1 face_data = pickle.loads(open("face_encodings.pickle", "rb").read()) def check_keys(model, pretrained_state_dict): ckpt_keys = set(pretrained_state_dict.keys()) model_keys = set(model.state_dict().keys()) used_pretrained_keys = model_keys & ckpt_keys unused_pretrained_keys = ckpt_keys - model_keys
def main(): args = get_args() torch.set_grad_enabled(False) cfg = None if args.network == "mobile0.25": cfg = cfg_mnet elif args.network == "resnet50": cfg = cfg_re50 # net and model net = RetinaFace(cfg=cfg, phase="test") net = load_model(net, args.trained_model, args.cpu) net.eval() print("Finished loading model!") print(net) cudnn.benchmark = True device = torch.device("cpu" if args.cpu else "cuda") net = net.to(device) args.save_folder.mkdir(exist_ok=True) fw = open(os.path.join(args.save_folder, args.dataset + "_dets.txt"), "w") # testing dataset testset_folder = os.path.join("data", args.dataset, "images/") testset_list = os.path.join("data", args.dataset, "img_list.txt") with open(testset_list, "r") as fr: test_dataset = fr.read().split() num_images = len(test_dataset) # testing scale resize = 1 _t = {"forward_pass": Timer(), "misc": Timer()} # testing begin for i, img_name in enumerate(test_dataset): image_path = testset_folder + img_name + ".jpg" img_raw = cv2.imread(image_path, cv2.IMREAD_COLOR) img = np.float32(img_raw) if resize != 1: img = cv2.resize(img, None, None, fx=resize, fy=resize, interpolation=cv2.INTER_LINEAR) im_height, im_width, _ = img.shape 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) _t["forward_pass"].tic() loc, conf, landms = net(img) # forward pass _t["forward_pass"].toc() _t["misc"].tic() priorbox = PriorBox(cfg, image_size=(im_height, im_width)) priors = priorbox.forward() priors = priors.to(device) prior_data = priors.data boxes = decode(loc.data.squeeze(0), prior_data, cfg["variance"]) boxes = boxes * scale / resize boxes = boxes.cpu().numpy() scores = conf.squeeze(0).data.cpu().numpy()[:, 1] landms = decode_landm(landms.data.squeeze(0), prior_data, cfg["variance"]) scale1 = torch.Tensor([ img.shape[3], img.shape[2], img.shape[3], img.shape[2], img.shape[3], img.shape[2], img.shape[3], img.shape[2], img.shape[3], img.shape[2], ]) scale1 = scale1.to(device) landms = landms * scale1 / resize landms = landms.cpu().numpy() # ignore low scores inds = np.where(scores > args.confidence_threshold)[0] boxes = boxes[inds] landms = landms[inds] scores = scores[inds] # keep top-K before NMS # order = scores.argsort()[::-1][:args.top_k] order = scores.argsort()[::-1] boxes = boxes[order] landms = landms[order] scores = scores[order] # do NMS dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False) keep = py_cpu_nms(dets, args.nms_threshold) dets = dets[keep, :] landms = landms[keep] # keep top-K faster NMS # dets = dets[:args.keep_top_k, :] # landms = landms[:args.keep_top_k, :] dets = np.concatenate((dets, landms), axis=1) _t["misc"].toc() # save dets if args.dataset == "FDDB": fw.write("{:s}\n".format(img_name)) fw.write("{:.1f}\n".format(dets.shape[0])) for k in range(dets.shape[0]): xmin = dets[k, 0] ymin = dets[k, 1] xmax = dets[k, 2] ymax = dets[k, 3] score = dets[k, 4] w = xmax - xmin + 1 h = ymax - ymin + 1 # fw.write('{:.3f} {:.3f} {:.3f} {:.3f} {:.10f}\n'.format(xmin, ymin, w, h, score)) fw.write("{:d} {:d} {:d} {:d} {:.10f}\n".format( int(xmin), int(ymin), int(w), int(h), score)) print("im_detect: {:d}/{:d} forward_pass_time: {:.4f}s misc: {:.4f}s". format(i + 1, num_images, _t["forward_pass"].average_time, _t["misc"].average_time)) # show image if args.save_image: for b in dets: if b[4] < args.vis_thres: continue text = "{:.4f}".format(b[4]) b = list(map(int, b)) cv2.rectangle(img_raw, (b[0], b[1]), (b[2], b[3]), (0, 0, 255), 2) cx = b[0] cy = b[1] + 12 cv2.putText(img_raw, text, (cx, cy), cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255)) # landms cv2.circle(img_raw, (b[5], b[6]), 1, (0, 0, 255), 4) cv2.circle(img_raw, (b[7], b[8]), 1, (0, 255, 255), 4) cv2.circle(img_raw, (b[9], b[10]), 1, (255, 0, 255), 4) cv2.circle(img_raw, (b[11], b[12]), 1, (0, 255, 0), 4) cv2.circle(img_raw, (b[13], b[14]), 1, (255, 0, 0), 4) # save image if not os.path.exists("./results/"): os.makedirs("./results/") name = "./results/" + str(i) + ".jpg" cv2.imwrite(name, img_raw) fw.close()