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
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
net = None
if args.network == "mobile0.25":
cfg = cfg_mnet
net = RetinaFace(cfg = cfg, phase = 'test')
elif args.network == "slim":
cfg = cfg_slim
net = Slim(cfg = cfg, phase = 'test')
elif args.network == "RFB":
cfg = cfg_rfb
net = RFB(cfg = cfg, phase = 'test')
else:
print("Don't support network!")
exit(0)
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)
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)
rgb_mean = (104, 117, 123) # bgr order
num_classes = 2