def init_detection_model(model_name, half=False, device='cuda'):
if model_name == 'retinaface_resnet50':
model = RetinaFace(network_name='resnet50', half=half)
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth'
elif model_name == 'retinaface_mobile0.25':
model = RetinaFace(network_name='mobile0.25', half=half)
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_mobilenet0.25_Final.pth'
else:
raise NotImplementedError(f'{model_name} is not implemented.')
model_path = load_file_from_url(url=model_url,
model_dir='facexlib/weights',
progress=True,
file_name=None)
# TODO: clean pretrained model
load_net = torch.load(model_path,
map_location=lambda storage, loc: storage)
# remove unnecessary 'module.'
for k, v in deepcopy(load_net).items():
if k.startswith('module.'):
load_net[k[7:]] = v
load_net.pop(k)
model.load_state_dict(load_net, strict=True)
model.eval()
model = model.to(device)
return model
def init_alignment_model(model_name, half=False, device='cuda'):
if model_name == 'awing_fan':
model = FAN(num_modules=4, num_landmarks=98)
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/alignment_WFLW_4HG.pth'
else:
raise NotImplementedError(f'{model_name} is not implemented.')
model_path = load_file_from_url(url=model_url, model_dir='facexlib/weights', progress=True, file_name=None)
model.load_state_dict(torch.load(model_path)['state_dict'], strict=True)
model.eval()
model = model.to(device)
return model
def init_headpose_model(model_name, half=False, device='cuda'):
if model_name == 'hopenet':
model = HopeNet('resnet', [3, 4, 6, 3], 66)
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.2.0/headpose_hopenet.pth'
else:
raise NotImplementedError(f'{model_name} is not implemented.')
model_path = load_file_from_url(url=model_url, model_dir='facexlib/weights', progress=True, file_name=None)
load_net = torch.load(model_path, map_location=lambda storage, loc: storage)['params']
model.load_state_dict(load_net, strict=True)
model.eval()
model = model.to(device)
return model
def init_recognition_model(model_name, half=False, device='cuda'):
if model_name == 'arcface':
model = Backbone(num_layers=50, drop_ratio=0.6,
mode='ir_se').to('cuda').eval()
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.1.0/recognition_arcface_ir_se50.pth'
else:
raise NotImplementedError(f'{model_name} is not implemented.')
model_path = load_file_from_url(url=model_url,
model_dir='facexlib/weights',
progress=True,
file_name=None)
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
return model
def init_assessment_model(model_name, half=False, device='cuda'):
if model_name == 'hypernet':
model = HyperIQA(16, 112, 224, 112, 56, 28, 14, 7)
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.2.0/assessment_hyperIQA.pth'
else:
raise NotImplementedError(f'{model_name} is not implemented.')
# load the pre-trained hypernet model
hypernet_model_path = load_file_from_url(url=model_url,
model_dir='facexlib/weights',
progress=True,
file_name=None)
model.hypernet.load_state_dict(
(torch.load(hypernet_model_path,
map_location=lambda storage, loc: storage)))
model = model.eval()
model = model.to(device)
return model
def init_matting_model(model_name='modnet', half=False, device='cuda'):
if model_name == 'modnet':
model = MODNet(backbone_pretrained=False)
model_url = 'https://github.com/xinntao/facexlib/releases/download/v0.2.0/matting_modnet_portrait.pth'
else:
raise NotImplementedError(f'{model_name} is not implemented.')
model_path = load_file_from_url(url=model_url,
model_dir='facexlib/weights',
progress=True,
file_name=None)
# TODO: clean pretrained model
load_net = torch.load(model_path,
map_location=lambda storage, loc: storage)
# remove unnecessary 'module.'
for k, v in deepcopy(load_net).items():
if k.startswith('module.'):
load_net[k[7:]] = v
load_net.pop(k)
model.load_state_dict(load_net, strict=True)
model.eval()
model = model.to(device)
return model