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()
