Spaces:
Running
Running
| """ | |
| =============================================================================== | |
| Author: Anjith George | |
| Institution: Idiap Research Institute, Martigny, Switzerland. | |
| Copyright (C) 2024 Anjith George | |
| This software is distributed under the terms described in the LICENSE file | |
| located in the parent directory of this source code repository. | |
| For inquiries, please contact the author at [email protected] | |
| =============================================================================== | |
| """ | |
| dependencies = ['torch', 'torchvision', 'timm'] | |
| from backbones import get_model | |
| import torch | |
| def edgeface_base(pretrained=True, **kwargs): | |
| model = get_model('edgeface_base', **kwargs) | |
| if pretrained: | |
| checkpoint_url = 'https://gitlab.idiap.ch/bob/bob.paper.tbiom2023_edgeface/-/raw/master/checkpoints/edgeface_base.pt' | |
| state_dict = torch.hub.load_state_dict_from_url( | |
| checkpoint_url, map_location='cpu' | |
| ) | |
| model.load_state_dict(state_dict) | |
| return model | |
| def edgeface_xs_gamma_06(pretrained=True, **kwargs): | |
| model = get_model('edgeface_xs_gamma_06', **kwargs) | |
| if pretrained: | |
| checkpoint_url = 'https://gitlab.idiap.ch/bob/bob.paper.tbiom2023_edgeface/-/raw/master/checkpoints/edgeface_xs_gamma_06.pt' | |
| state_dict = torch.hub.load_state_dict_from_url( | |
| checkpoint_url, map_location='cpu' | |
| ) | |
| model.load_state_dict(state_dict) | |
| return model | |
| def edgeface_xs_q(pretrained=True, **kwargs): | |
| model = get_model('edgeface_xs_q', **kwargs) | |
| if pretrained: | |
| checkpoint_url = 'https://gitlab.idiap.ch/bob/bob.paper.tbiom2023_edgeface/-/raw/master/checkpoints/edgeface_xs_q.pt' | |
| state_dict = torch.hub.load_state_dict_from_url( | |
| checkpoint_url, map_location='cpu' | |
| ) | |
| model.load_state_dict(state_dict) | |
| return model | |
| def edgeface_xxs(pretrained=True, **kwargs): | |
| model = get_model('edgeface_xxs', **kwargs) | |
| if pretrained: | |
| checkpoint_url = 'https://gitlab.idiap.ch/bob/bob.paper.tbiom2023_edgeface/-/raw/master/checkpoints/edgeface_xxs.pt' | |
| state_dict = torch.hub.load_state_dict_from_url( | |
| checkpoint_url, map_location='cpu' | |
| ) | |
| model.load_state_dict(state_dict) | |
| return model | |
| def edgeface_xxs_q(pretrained=True, **kwargs): | |
| model = get_model('edgeface_xxs_q', **kwargs) | |
| if pretrained: | |
| checkpoint_url = 'https://gitlab.idiap.ch/bob/bob.paper.tbiom2023_edgeface/-/raw/master/checkpoints/edgeface_xxs_q.pt' | |
| state_dict = torch.hub.load_state_dict_from_url( | |
| checkpoint_url, map_location='cpu' | |
| ) | |
| model.load_state_dict(state_dict) | |
| return model | |
| def edgeface_s_gamma_05(pretrained=True, **kwargs): | |
| model = get_model('edgeface_s_gamma_05', **kwargs) | |
| if pretrained: | |
| checkpoint_url = 'https://gitlab.idiap.ch/bob/bob.paper.tbiom2023_edgeface/-/raw/master/checkpoints/edgeface_s_gamma_05.pt' | |
| state_dict = torch.hub.load_state_dict_from_url( | |
| checkpoint_url, map_location='cpu' | |
| ) | |
| model.load_state_dict(state_dict) | |
| return model | |