mirror of
https://github.com/yakhyo/uniface.git
synced 2026-05-20 08:25:08 +00:00
9bf54f5f78
* refactor: Split recognition models into separate files * feat: Add EdgeFace recognition model * release: Bump version to v3.4.0
67 lines
1.6 KiB
Python
67 lines
1.6 KiB
Python
import argparse
|
|
|
|
from uniface.constants import (
|
|
AdaFaceWeights,
|
|
AgeGenderWeights,
|
|
ArcFaceWeights,
|
|
DDAMFNWeights,
|
|
EdgeFaceWeights,
|
|
HeadPoseWeights,
|
|
LandmarkWeights,
|
|
MobileFaceWeights,
|
|
RetinaFaceWeights,
|
|
SCRFDWeights,
|
|
SphereFaceWeights,
|
|
)
|
|
from uniface.model_store import verify_model_weights
|
|
|
|
MODEL_TYPES = {
|
|
'retinaface': RetinaFaceWeights,
|
|
'adaface': AdaFaceWeights,
|
|
'arcface': ArcFaceWeights,
|
|
'edgeface': EdgeFaceWeights,
|
|
'mobileface': MobileFaceWeights,
|
|
'sphereface': SphereFaceWeights,
|
|
'scrfd': SCRFDWeights,
|
|
'ddamfn': DDAMFNWeights,
|
|
'agegender': AgeGenderWeights,
|
|
'landmark': LandmarkWeights,
|
|
'headpose': HeadPoseWeights,
|
|
}
|
|
|
|
|
|
def download_models(model_enum):
|
|
for weight in model_enum:
|
|
print(f'Downloading: {weight.value}')
|
|
try:
|
|
verify_model_weights(weight)
|
|
print(f' Done: {weight.value}')
|
|
except Exception as e:
|
|
print(f' Failed: {e}')
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Download model weights')
|
|
parser.add_argument(
|
|
'--model-type',
|
|
type=str,
|
|
choices=list(MODEL_TYPES.keys()),
|
|
help='Model type to download. If not specified, downloads all.',
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if args.model_type:
|
|
print(f'Downloading {args.model_type} models...')
|
|
download_models(MODEL_TYPES[args.model_type])
|
|
else:
|
|
print('Downloading all models...')
|
|
for name, model_enum in MODEL_TYPES.items():
|
|
print(f'\n{name}:')
|
|
download_models(model_enum)
|
|
|
|
print('\nDone!')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|