INNER CODE UNIT · Python
__init__
minivision-ai/Silent-Face-Anti-Spoofing · src/anti_spoof_predict.py:54
def __init__(self, device_id):
super(AntiSpoofPredict, self).__init__()
self.device = torch.device("cuda:{}".format(device_id)
if torch.cuda.is_available() else "cpu")
def _load_model(self, model_path):
# define model
model_name = os.path.basename(model_path)
h_input, w_input, model_type, _ = parse_model_name(model_name)
self.kernel_size = get_kernel(h_input, w_input,)
self.model = MODEL_MAPPING[model_type](conv6_kernel=self.kernel_size).to(self.device)
# load model weight
state_dict = torch.load(model_path, map_location=self.device)
keys = iter(state_dict)
first_layer_name = keys.__next__()
if first_layer_name.find('module.') >= 0:
from collections import OrderedDict