def create_XNet(device, p_model_path, r_model_path, o_model_path):
m_PNet, m_RNet, m_ONet = None, None, None
if os.access(p_model_path, os.F_OK): # PNet
m_PNet = PNet()
m_PNet = m_PNet.to(device)
chkpt = torch.load(p_model_path, map_location=device)
m_PNet.load_state_dict(chkpt)
m_PNet.eval()
print('load PNet ~ : ', p_model_path)
if os.access(r_model_path, os.F_OK): # RNet
# m_RNet.load_state_dict(torch.load(r_model_path))
m_RNet = RNet()
m_RNet = m_RNet.to(device)
chkpt = torch.load(r_model_path, map_location=device)
m_RNet.load_state_dict(chkpt)
m_RNet.eval()
print('load RNet ~ : ', r_model_path)
if os.access(o_model_path, os.F_OK): # RNet
# m_RNet.load_state_dict(torch.load(r_model_path))
m_ONet = ONet()
m_ONet = m_ONet.to(device)
chkpt = torch.load(o_model_path, map_location=device)
m_ONet.load_state_dict(chkpt)
m_ONet.eval()
print('load ONet ~ : ', o_model_path)
return m_PNet, m_RNet, m_ONet
def create_mtcnn_net(p_model_path=None,
r_model_path=None,
o_model_path=None,
use_cuda=True):
pnet, rnet, onet = None, None, None
if p_model_path is not None:
pnet = PNet(use_cuda=use_cuda)
if (use_cuda):
print('p_model_path:{0}'.format(p_model_path))
pnet.load_state_dict(torch.load(p_model_path))
pnet.cuda()
else:
# forcing all GPU tensors to be in CPU while loading
pnet.load_state_dict(
torch.load(p_model_path,
map_location=lambda storage, loc: storage))
pnet.eval()
if r_model_path is not None:
rnet = RNet(use_cuda=use_cuda)
if (use_cuda):
print('r_model_path:{0}'.format(r_model_path))
rnet.load_state_dict(torch.load(r_model_path))
rnet.cuda()
else:
rnet.load_state_dict(
torch.load(r_model_path,
map_location=lambda storage, loc: storage))
rnet.eval()
if o_model_path is not None:
onet = ONet(use_cuda=use_cuda)
if (use_cuda):
print('o_model_path:{0}'.format(o_model_path))
onet.load_state_dict(torch.load(o_model_path))
onet.cuda()
else:
onet.load_state_dict(
torch.load(o_model_path,
map_location=lambda storage, loc: storage))
onet.eval()
return pnet, rnet, onet
def create_mtcnn_model(p_model_path=None,
r_model_path=None,
o_model_path=None):
pnet, rnet, onet = None, None, None
if p_model_path is not None:
pnet = PNet()
pnet.load_state_dict(torch.load(p_model_path))
pnet.to('cuda')
pnet.eval()
# pnet.load_state_dict(torch.load(p_model_path, map_location=lambda storage, loc: storage))
if r_model_path is not None:
rnet = RNet()
rnet.load_state_dict(torch.load(r_model_path))
rnet.to('cuda')
rnet.eval()
if o_model_path is not None:
onet = ONet()
onet.load_state_dict(torch.load(o_model_path))
onet.to('cuda')
onet.eval()
return pnet, rnet, onet