| """ |
| This code file mainly comes from https://github.com/dmlc/gluon-cv/blob/master/gluoncv/model_zoo/model_store.py |
| """ |
| from __future__ import print_function |
|
|
| __all__ = ['get_model_file'] |
| import os |
| import zipfile |
| import glob |
|
|
| from ..utils import download, check_sha1 |
|
|
| _model_sha1 = { |
| name: checksum |
| for checksum, name in [ |
| ('95be21b58e29e9c1237f229dae534bd854009ce0', 'arcface_r100_v1'), |
| ('', 'arcface_mfn_v1'), |
| ('39fd1e087a2a2ed70a154ac01fecaa86c315d01b', 'retinaface_r50_v1'), |
| ('2c9de8116d1f448fd1d4661f90308faae34c990a', 'retinaface_mnet025_v1'), |
| ('0db1d07921d005e6c9a5b38e059452fc5645e5a4', 'retinaface_mnet025_v2'), |
| ('7dd8111652b7aac2490c5dcddeb268e53ac643e6', 'genderage_v1'), |
| ] |
| } |
|
|
| base_repo_url = 'https://insightface.ai/files/' |
| _url_format = '{repo_url}models/{file_name}.zip' |
|
|
|
|
| def short_hash(name): |
| if name not in _model_sha1: |
| raise ValueError( |
| 'Pretrained model for {name} is not available.'.format(name=name)) |
| return _model_sha1[name][:8] |
|
|
|
|
| def find_params_file(dir_path): |
| if not os.path.exists(dir_path): |
| return None |
| paths = glob.glob("%s/*.params" % dir_path) |
| if len(paths) == 0: |
| return None |
| paths = sorted(paths) |
| return paths[-1] |
|
|
|
|
| def get_model_file(name, root=os.path.join('~', '.insightface', 'models')): |
| r"""Return location for the pretrained on local file system. |
| |
| This function will download from online model zoo when model cannot be found or has mismatch. |
| The root directory will be created if it doesn't exist. |
| |
| Parameters |
| ---------- |
| name : str |
| Name of the model. |
| root : str, default '~/.mxnet/models' |
| Location for keeping the model parameters. |
| |
| Returns |
| ------- |
| file_path |
| Path to the requested pretrained model file. |
| """ |
|
|
| file_name = name |
| root = os.path.expanduser(root) |
| dir_path = os.path.join(root, name) |
| file_path = find_params_file(dir_path) |
| |
| sha1_hash = _model_sha1[name] |
| if file_path is not None: |
| if check_sha1(file_path, sha1_hash): |
| return file_path |
| else: |
| print( |
| 'Mismatch in the content of model file detected. Downloading again.' |
| ) |
| else: |
| print('Model file is not found. Downloading.') |
|
|
| if not os.path.exists(root): |
| os.makedirs(root) |
| if not os.path.exists(dir_path): |
| os.makedirs(dir_path) |
|
|
| zip_file_path = os.path.join(root, file_name + '.zip') |
| repo_url = base_repo_url |
| if repo_url[-1] != '/': |
| repo_url = repo_url + '/' |
| download(_url_format.format(repo_url=repo_url, file_name=file_name), |
| path=zip_file_path, |
| overwrite=True) |
| with zipfile.ZipFile(zip_file_path) as zf: |
| zf.extractall(dir_path) |
| os.remove(zip_file_path) |
| file_path = find_params_file(dir_path) |
|
|
| if check_sha1(file_path, sha1_hash): |
| return file_path |
| else: |
| raise ValueError( |
| 'Downloaded file has different hash. Please try again.') |
|
|
|
|
