-
Notifications
You must be signed in to change notification settings - Fork 39
/
generate_masks_database.py
57 lines (42 loc) · 1.19 KB
/
generate_masks_database.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
from argparse import ArgumentParser
from masked_face_sdk.mask_generation_utils import generate_masks_base
import face_alignment
import json
def parse_args():
parser = ArgumentParser(description='Generate masks database')
parser.add_argument(
'--masks-folder', required=True, type=str,
help='Path to folder with masks images.'
)
parser.add_argument(
'--database-file', required=True, type=str,
help='Path to created json database file.'
)
parser.add_argument(
'--verbose', action='store_true'
)
parser.add_argument(
'--skip-warnings', action='store_true'
)
return parser.parse_args()
def main():
args = parse_args()
fa = face_alignment.FaceAlignment(
face_alignment.LandmarksType._2D,
device='cpu'
)
data_dict = generate_masks_base(
args.masks_folder,
fa,
args.verbose,
args.skip_warnings
)
with open(args.database_file, 'w') as jf:
json.dump(data_dict, jf, indent=4)
print(
'Masks database successfully saved by follow path: {}'.format(
args.database_file
)
)
if __name__ == '__main__':
main()