/
enroll-mmu2.py
68 lines (53 loc) · 2.4 KB
/
enroll-mmu2.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
58
59
60
61
62
63
64
65
66
67
68
# -----------------------------------------------------------------------------
# Import
# -----------------------------------------------------------------------------
import argparse
import os
from glob import glob
from tqdm import tqdm
from time import time
from scipy.io import savemat
from multiprocessing import cpu_count, Pool
from fnc.extractFeature import extractFeature
# ------------------------------------------------------------------------------
# Argument parsing
# ------------------------------------------------------------------------------
parser = argparse.ArgumentParser()
parser.add_argument("--data_dir", type=str, default="../MMU2/",
help="Path to the directory containing MMU2 images.")
parser.add_argument("--temp_dir", type=str, default="./templates/MMU2/",
help="Path to the directory containing templates.")
parser.add_argument("--n_cores", type=int, default=cpu_count(),
help="Number of cores used for enrolling template.")
args = parser.parse_args()
# -----------------------------------------------------------------------------
# Pool function
# -----------------------------------------------------------------------------
def pool_func(file):
template, mask, _ = extractFeature(file, use_multiprocess=False)
basename = os.path.basename(file)
out_file = os.path.join(args.temp_dir, "%s.mat" % (basename))
savemat(out_file, mdict={'template': template, 'mask': mask})
# -----------------------------------------------------------------------------
# Execution
# -----------------------------------------------------------------------------
start = time()
# Check the existence of temp_dir
if not os.path.exists(args.temp_dir):
print("makedirs", args.temp_dir)
os.makedirs(args.temp_dir)
# Get list of files for enrolling template.
# Just "*010*.jpg" files are selected.
# In addition, 50th samples are not used.
files = glob(os.path.join(args.data_dir, "**/*/*010*.bmp"), recursive=True)
files = [file for file in files if "01020" not in file]
files = [file for file in files if "5002" not in file]
n_files = len(files)
print("Number of files for enrolling:", n_files)
# Parallel pools to enroll templates
print("Start enrolling...")
pools = Pool(processes=args.n_cores)
for _ in tqdm(pools.imap_unordered(pool_func, files), total=n_files):
pass
end = time()
print('\n>>> Enrollment time: {} [s]\n'.format(end-start))