├── .gitignore
├── ERC
├── .DS_Store
├── README.md
├── erc.py
├── erc_ijbc.py
└── roc.py
├── IJB-C Evaluation
├── Readme.md
├── eval_ijbc_qs.py
├── extract_IJB.py
├── extract_emb.py
└── ijb_qs_pair_file.py
├── README.md
├── backbones
├── __init__.py
└── iresnet.py
├── config.py
├── dataset.py
├── eval
└── verification.py
├── evaluation
├── FaceModel.py
├── QualityModel.py
├── crerate_pair_xqlfw.py
├── extract_adience.py
└── getQualityScore.py
├── feature_extraction
├── backbones
│ ├── activation.py
│ ├── iresnet.py
│ ├── iresnet_mag.py
│ ├── mag_network_inf.py
│ ├── model_irse.py
│ └── utils.py
├── eval_ijbc_qs.py
├── extract_IJB.py
├── extract_adience.py
├── extract_bin.py
├── extract_emb.py
├── extract_xqlfw.py
├── ijb_pair_file.py
└── model
│ ├── ArcFaceModel.py
│ ├── CurricularFaceModel.py
│ ├── ElasticFaceModel.py
│ ├── FaceModel.py
│ └── MagFaceModel.py
├── losses.py
├── requirement.txt
├── run.sh
├── train_cr_fiqa.py
└── utils
├── NIST.png
├── __init__.py
├── align_trans.py
├── utils_amp.py
├── utils_callbacks.py
├── utils_logging.py
└── workflow.png
/.gitignore:
--------------------------------------------------------------------------------
1 | # Byte-compiled / optimized / DLL files
2 | __pycache__/
3 | *.py[cod]
4 | *$py.class
5 |
6 | # C extensions
7 | *.so
8 |
9 | # Distribution / packaging
10 | .Python
11 | build/
12 | develop-eggs/
13 | dist/
14 | downloads/
15 | eggs/
16 | .eggs/
17 | lib/
18 | lib64/
19 | parts/
20 | sdist/
21 | var/
22 | wheels/
23 | pip-wheel-metadata/
24 | share/python-wheels/
25 | *.egg-info/
26 | .installed.cfg
27 | *.egg
28 | MANIFEST
29 |
30 | # PyInstaller
31 | # Usually these files are written by a python script from a template
32 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
33 | *.manifest
34 | *.spec
35 |
36 | # Installer logs
37 | pip-log.txt
38 | pip-delete-this-directory.txt
39 |
40 | # Unit test / coverage reports
41 | htmlcov/
42 | .tox/
43 | .nox/
44 | .coverage
45 | .coverage.*
46 | .cache
47 | nosetests.xml
48 | coverage.xml
49 | *.cover
50 | *.py,cover
51 | .hypothesis/
52 | .pytest_cache/
53 |
54 | # Translations
55 | *.mo
56 | *.pot
57 |
58 | # Django stuff:
59 | *.log
60 | local_settings.py
61 | db.sqlite3
62 | db.sqlite3-journal
63 |
64 | # Flask stuff:
65 | instance/
66 | .webassets-cache
67 |
68 | # Scrapy stuff:
69 | .scrapy
70 |
71 | # Sphinx documentation
72 | docs/_build/
73 |
74 | # PyBuilder
75 | target/
76 |
77 | # Jupyter Notebook
78 | .ipynb_checkpoints
79 |
80 | # IPython
81 | profile_default/
82 | ipython_config.py
83 |
84 | # pyenv
85 | .python-version
86 |
87 | # pipenv
88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not
91 | # install all needed dependencies.
92 | #Pipfile.lock
93 |
94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow
95 | __pypackages__/
96 |
97 | # Celery stuff
98 | celerybeat-schedule
99 | celerybeat.pid
100 |
101 | # SageMath parsed files
102 | *.sage.py
103 |
104 | # Environments
105 | .env
106 | .venv
107 | env/
108 | venv/
109 | ENV/
110 | env.bak/
111 | venv.bak/
112 |
113 | # Spyder project settings
114 | .spyderproject
115 | .spyproject
116 |
117 | # Rope project settings
118 | .ropeproject
119 |
120 | # mkdocs documentation
121 | /site
122 |
123 | # mypy
124 | .mypy_cache/
125 | .dmypy.json
126 | dmypy.json
127 |
128 | # Pyre type checker
129 | .pyre/
130 |
131 | .DS_Store
132 | .idea
--------------------------------------------------------------------------------
/ERC/.DS_Store:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/fdbtrs/CR-FIQA/a8f98da6ee03db1e8dc091b5910931d41d69fa44/ERC/.DS_Store
--------------------------------------------------------------------------------
/ERC/README.md:
--------------------------------------------------------------------------------
1 | # ERC plot
2 | ---
3 |
4 | ## Data preparation
5 | For plotting Error vs. Reject Curve (ERC) and using the following code, the format of the directories and files are:
6 | 1. Embedding directory that save embeddings computed from different FR models:
7 | ```
8 | +-- embeddings_dir
9 | | +-- dataset_FRmodel #e.g., lfw_ArcFaceModel
10 | | +-- xxx.npy
11 | ```
12 | 2. #### Quality score directory that save image pair list and quality scores
13 | ```
14 | +-- quality_score_dir
15 | | +-- dataset # e.g., lfw
16 | | +-- images # where save the face images
17 | | +-- pair_list.txt
18 | | +-- qualityEstimationMethod_dataset.txt #e.g., CRFIQAL_lfw.txt
19 | ```
20 | where the format of pair_list.txt is:
21 | ```
22 | image_name_1, image_name_2, type
23 | # 1: genuine, 0: Imposter
24 | 00000 00001 1
25 | 00002 00003 1
26 | 05402 05405 0
27 | 05404 05409 0
28 | ```
29 | and the format of qualityEstimationMethod_dataset.txt is:
30 | ```
31 | image_path, quality score
32 | quality_score_dir/lfw/images/00000.jpg 0.36241477727890015
33 | quality_score_dir/lfw/images/00000.jpg 0.36241477727890015
34 | quality_score_dir/lfw/images/00001.jpg 0.37981975078582764
35 | quality_score_dir/lfw/images/00002.jpg 0.44192782044410706
36 | quality_score_dir/lfw/images/00003.jpg 0.5173501372337341
37 | ```
38 |
39 | ## ERC evaluation
40 | 1. After preparing the above data, erc.py can be used for compute and plot:
41 | ```
42 | python erc.py \
43 | --embeddings_dir 'embeddings_dir' \
44 | --quality_score_dir 'quality_score_dir' \
45 | --method_name 'CRFIQAL,CRFIQAS,FaceQnet,DeepIQA,PFE,..' \
46 | --models 'ArcFaceModel, ElasticFaceModel, MagFaceModel, CurricularFaceModel' \
47 | --eval_db 'adience,XQLFW,lfw,agedb_30,calfw,cplfw,cfp_fp' \
48 | --distance_metric 'cosine' \
49 | --output_dir 'output_dir' \
50 | ```
51 | More details can be checked in erc.py
52 |
53 | ---
54 |
55 | ## Data preparation for IJB-C
56 | Becuase the IJB-C dataset is a large-scale dataset, the way for loading embeddings and quality scores is a little bit different due to the limited computation resource.
57 | For plotting Error vs. Reject Curve (ERC) and using the following code, the format of the directories and files are:
58 | 1. Embedding directory that save embeddings computed from different FR models:
59 | ```
60 | +-- embeddings_dir
61 | | +-- dataset_FRmodel #e.g., IJBC_ArcFaceModel
62 | | +-- xxx.npy
63 | ```
64 | 2. Quality score directory that save image pair list and quality scores
65 | ```
66 | +-- quality_score_dir
67 | | +-- dataset # e.g., lfw
68 | | +-- images # where save the face images
69 | | +-- qualityEstimationMethod_dataset.txt #e.g., CRFIQAL_IJBC.txt
70 | ```
71 | where the format of qualityEstimationMethod_dataset.txt is:
72 | ```
73 | image_name_1, image_name_2, type, quality score # in type column: 1: genuine, 0: Imposter
74 | 171707 187569 0 1.7933012
75 | 26 17991 0 1.764878
76 | 878 20568 1 0.5316349
77 | 170001 178303 1 0.10600686
78 | ```
79 |
80 | ## ERC evaluation for IJB-C
81 | 1. After preparing the above data, erc_ijbc.py can be used for compute and plot:
82 | ```
83 | python erc.py \
84 | --embeddings_dir 'embeddings_dir' \
85 | --quality_score_dir 'quality_score_dir' \
86 | --method_name 'CRFIQAL,CRFIQAS,FaceQnet,DeepIQA,PFE,..' \
87 | --models 'ArcFaceModel, ElasticFaceModel, MagFaceModel, CurricularFaceModel' \
88 | --eval_db 'IJBC' \
89 | --distance_metric 'cosine' \
90 | --output_dir 'output_dir' \
91 | ```
92 | More details can be checked in erc_ijbc.py
93 |
--------------------------------------------------------------------------------
/ERC/erc.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | import argparse
3 | import os
4 | import glob
5 | import math
6 | import numpy as np
7 | from tqdm import tqdm
8 | import sklearn
9 | from sklearn import metrics
10 |
11 | from collections import defaultdict
12 | import gc
13 |
14 | import matplotlib.pyplot as plt
15 | from matplotlib.font_manager import FontProperties
16 |
17 | from ERC.roc import get_eer_threshold
18 |
19 | parser = argparse.ArgumentParser(description='Evaluation')
20 |
21 | parser.add_argument('--embeddings_dir', type=str,
22 | default="/usr/quality_embeddings",
23 | help='The dir save embeddings for each method and dataset, the diretory inside should be: {dataset}_{model}, e.g., IJBC_ArcFaceModel')
24 | parser.add_argument('--quality_score_dir', type=str,
25 | default="/usr/quality_score",
26 | help='The dir save file of quality scores for each dataset and method, the file inside should be: {method}_{dataset}.txt, e.g., CRFIQAS_IJBC.txt')
27 |
28 | parser.add_argument('--method_name', type=str,
29 | default="Serfiq,BRISQUE,SDD,rankIQ,magface,FaceQnet,DeepIQA,PFE,rankIQA,CRFIQAL,CRFIQAS",
30 | help='The evaluated image quality estimation method')
31 | parser.add_argument('--models', type=str,
32 | default="ArcFaceModel, ElasticFaceModel, MagFaceModel, CurricularFaceModel",
33 | help='The evaluated FR model')
34 | parser.add_argument('--eval_db', type=str,
35 | default="adience,XQLFW,lfw,agedb_30,calfw,cplfw,cfp_fp",
36 | help='The evaluated dataset')
37 |
38 | parser.add_argument('--distance_metric', type=str,
39 | default='cosine',
40 | help='Cosine distance or euclidian distance')
41 | parser.add_argument('--output_dir', type=str,
42 | default="erc_plot_auc_test_all",
43 | help='')
44 |
45 | IMAGE_EXTENSION = '.jpg'
46 |
47 | def load_quality(scores):
48 | quality={}
49 | with open(scores[0], 'r') as f:
50 | lines=f.readlines()
51 | for l in lines:
52 | scores = l.split()[1].strip()
53 | n = l.split()[0].strip()
54 | quality[n] = scores
55 | return quality
56 |
57 | def load_quality_pair(pair_path, scores, dataset, args):
58 | pairs_quality = []
59 | quality=load_quality(scores)
60 | with open(pair_path, 'r') as f:
61 | lines=f.readlines()
62 | for idex in range(len(lines)):
63 | a= lines[idex].rstrip().split()[0]
64 | b= lines[idex].rstrip().split()[1]
65 | qlt=min(float(quality.get(os.path.join(args.quality_score_dir, dataset, 'images', f"{a}{IMAGE_EXTENSION}"))),
66 | float(quality.get(os.path.join(args.quality_score_dir, dataset, 'images', f"{b}{IMAGE_EXTENSION}"))))
67 | pairs_quality.append(qlt)
68 | return pairs_quality
69 |
70 | def load_feat_pair(pair_path, root):
71 | pairs = {}
72 | with open(pair_path, 'r') as f:
73 | lines=f.readlines()
74 | for idex in range(len(lines)):
75 | a= lines[idex].rstrip().split()[0]
76 | b= lines[idex].rstrip().split()[1]
77 | is_same=int(lines[idex].rstrip().split()[2])
78 | feat_a=np.load(os.path.join(root, f"{a}.npy"))
79 | feat_b=np.load(os.path.join(root, f"{b}.npy"))
80 | pairs[idex] = [feat_a, feat_b, is_same]
81 | print("All features are loaded")
82 | return pairs
83 |
84 | def distance_(embeddings0, embeddings1, dist="cosine"):
85 | # Distance based on cosine similarity
86 | if (dist=="cosine"):
87 | dot = np.sum(np.multiply(embeddings0, embeddings1), axis=1)
88 | norm = np.linalg.norm(embeddings0, axis=1) * np.linalg.norm(embeddings1, axis=1)
89 | # shaving
90 | similarity = np.clip(dot / norm, -1., 1.)
91 | dist = np.arccos(similarity) / math.pi
92 | else:
93 | embeddings0 = sklearn.preprocessing.normalize(embeddings0)
94 | embeddings1 = sklearn.preprocessing.normalize(embeddings1)
95 | diff = np.subtract(embeddings0, embeddings1)
96 | dist = np.sum(np.square(diff), 1)
97 |
98 | return dist
99 |
100 | def calc_score(embeddings0, embeddings1, actual_issame, subtract_mean=False, dist_type='cosine'):
101 | assert (embeddings0.shape[0] == embeddings1.shape[0])
102 | assert (embeddings0.shape[1] == embeddings1.shape[1])
103 |
104 | if subtract_mean:
105 | mean = np.mean(np.concatenate([embeddings0, embeddings1]), axis=0)
106 | else:
107 | mean = 0.
108 |
109 | dist = distance_(embeddings0, embeddings1, dist=dist_type)
110 | # sort in a desending order
111 | pos_scores =np.sort(dist[actual_issame == 1])
112 | neg_scores = np.sort(dist[actual_issame == 0])
113 | return pos_scores, neg_scores
114 |
115 | def save_pdf(fnmrs_lists, method_labels, model, output_dir, fmr, db):
116 | fontsize = 20
117 | colors = ['green', 'black', 'orange', 'plum', 'cyan', 'gold', 'gray', 'salmon', 'deepskyblue', 'red', 'blue',
118 | 'darkseagreen', 'seashell', 'hotpink', 'indigo', 'lightseagreen', 'khaki', 'brown', 'teal', 'darkcyan']
119 | STYLES = ['--', '-.', ':', 'v--', '^--', ',--', '<--', '>--', '1--',
120 | '-' ,'-' , '2--', '3--', '4--', '.--', 'p--', '*--', 'h--', 'H--', '+--', 'x--', 'd--', '|--', '---']
121 | unconsidered_rates = 100 * np.arange(0, 0.98, 0.05)
122 |
123 | fig, ax1 = plt.subplots() # added
124 | if (not os.path.isdir(output_dir)):
125 | os.makedirs(output_dir)
126 |
127 | for i in range(len(fnmrs_lists)):
128 | print(fnmrs_lists[i])
129 | plt.plot(unconsidered_rates[:len(fnmrs_lists[i])], fnmrs_lists[i], STYLES[i], color=colors[i],
130 | label=method_labels[i])
131 | auc_value = metrics.auc( np.array(unconsidered_rates[:len(fnmrs_lists[i])]/100), np.array(fnmrs_lists[i]))
132 | with open(os.path.join(output_dir, db, str(fmr)+"_auc.txt"), "a") as f:
133 | f.write(db + ':' + model + ':' + method_labels[i] + ':' + str(auc_value) + '\n')
134 | plt.xlabel('Ratio of unconsidered image [%]')
135 |
136 | plt.xlabel('Ratio of unconsidered image [%]', fontsize=fontsize)
137 | plt.xlim([0, 98])
138 | plt.xticks(np.arange(0, 98, 10), fontsize=fontsize)
139 | plt.title(f"Testing on {db}, FMR={fmr}" + f" ({model})", fontsize=fontsize) # update : -3
140 | plt.ylabel('FNMR', fontsize=fontsize)
141 |
142 | axbox = ax1.get_position()
143 | fig.legend(bbox_to_anchor=(axbox.x0 + 0.5 * axbox.width, axbox.y0 - 0.22), prop=FontProperties(size=12),
144 | loc='lower center', ncol=6)
145 | plt.tight_layout()
146 | plt.savefig(os.path.join(output_dir, db, db + '_' +str(fmr) +'_'+model + '.png'), bbox_inches='tight')
147 |
148 | def getFNMRFixedTH(feat_pairs, qlts, dist_type='cosine', desc=True):
149 | embeddings0, embeddings1, targets = [], [], []
150 | pair_qlt_list = [] # store the min qlt
151 | for k, v in feat_pairs.items():
152 | feat_a = v[0]
153 | feat_b = v[1]
154 | ab_is_same = int(v[2])
155 | # convert into np
156 | np_feat_a = np.asarray(feat_a, dtype=np.float64)
157 | np_feat_b = np.asarray(feat_b, dtype=np.float64)
158 | # append
159 | embeddings0.append(np_feat_a)
160 | embeddings1.append(np_feat_b)
161 | targets.append(ab_is_same)
162 |
163 | # evaluate
164 | embeddings0 = np.vstack(embeddings0)
165 | embeddings1 = np.vstack(embeddings1)
166 | targets = np.vstack(targets).reshape(-1, )
167 | qlts = np.array(qlts)
168 | if (desc):
169 | qlts_sorted_idx = np.argsort(qlts)
170 | else:
171 | qlts_sorted_idx = np.argsort(qlts)[::-1]
172 |
173 | num_pairs = len(targets)
174 | unconsidered_rates = np.arange(0, 0.98, 0.05)
175 |
176 | fnmrs_list_2 = []
177 | fnmrs_list_3 = []
178 | fnmrs_list_4 = []
179 | for u_rate in unconsidered_rates:
180 | hq_pairs_idx = qlts_sorted_idx[int(u_rate * num_pairs):]
181 | pos_dists, neg_dists = calc_score(embeddings0[hq_pairs_idx], embeddings1[hq_pairs_idx], targets[hq_pairs_idx], dist_type=dist_type)
182 | fmr100_th, fmr1000_th, fmr10000_th = get_eer_threshold(pos_dists, neg_dists, ds_scores=True)
183 |
184 | g_true = [g for g in pos_dists if g < fmr100_th]
185 | fnmrs_list_2.append(1- len(g_true)/(len(pos_dists)))
186 | g_true = [g for g in pos_dists if g < fmr1000_th]
187 | fnmrs_list_3.append(1 - len(g_true) / (len(pos_dists)))
188 | g_true = [g for g in pos_dists if g < fmr10000_th]
189 | fnmrs_list_4.append(1 - len(g_true) / (len(pos_dists)))
190 |
191 | return fnmrs_list_2,fnmrs_list_3,fnmrs_list_4,unconsidered_rates
192 |
193 |
194 | def perform_1v1_quality_eval(args):
195 | d = args.eval_db.split(',')
196 |
197 | for dataset in d:
198 | if os.path.exists(os.path.join(args.output_dir, dataset, str(1e-2)+"_auc.txt")):
199 | os.remove(os.path.join(args.output_dir, dataset, str(1e-2)+"_auc.txt"))
200 | if os.path.exists(os.path.join(args.output_dir, dataset, str(1e-3)+"_auc.txt")):
201 | os.remove(os.path.join(args.output_dir, dataset, str(1e-3)+"_auc.txt"))
202 | if os.path.exists(os.path.join(args.output_dir, dataset, str(1e-4)+"_auc.txt")):
203 | os.remove(os.path.join(args.output_dir, dataset, str(1e-4)+"_auc.txt"))
204 |
205 | models=args.models.split(',')
206 | for model in models:
207 | for dataset in d:
208 | method_names = args.method_name.split(',')
209 | fnmrs_list_2=[]
210 | fnmrs_list_3=[]
211 | fnmrs_list_4=[]
212 | method_labels=[]
213 |
214 | if (not os.path.isdir(os.path.join(args.output_dir, dataset, 'fnmr'))):
215 | os.makedirs(os.path.join(args.output_dir, dataset, 'fnmr'))
216 |
217 | unconsidered_rates = np.arange(0, 0.98, 0.05)
218 |
219 | for method_name in method_names:
220 | print(f"----process {model} {dataset} {method_name}-----------")
221 | desc = False if method_name == 'PFE' else True
222 |
223 | feat_pairs = load_feat_pair(os.path.join(args.quality_score_dir, dataset, 'pair_list.txt'),
224 | os.path.join(args.embeddings_dir, f"{dataset}_{model}"))
225 |
226 | quality_scores = load_quality_pair(os.path.join(args.quality_score_dir, dataset, 'pair_list.txt')
227 | [os.path.join(args.quality_score_dir, dataset, f"{method_name}_{dataset}.txt")],
228 | dataset, args)
229 |
230 | fnmr2, fnmr3, fnmr4, unconsidered_rates = getFNMRFixedTH(feat_pairs, quality_scores, dist_type=args.distance_metric, desc=desc)
231 | fnmrs_list_2.append(fnmr2)
232 | fnmrs_list_3.append(fnmr3)
233 | fnmrs_list_4.append(fnmr4)
234 | method_labels.append(f"{method_name}")
235 |
236 | np.save(os.path.join(args.output_dir, dataset, 'fnmr', f"{method_name}_{model}_{dataset}_fnmr2.npy"), fnmr2)
237 | np.save(os.path.join(args.output_dir, dataset, 'fnmr', f"{method_name}_{model}_{dataset}_fnmr3.npy"), fnmr3)
238 | np.save(os.path.join(args.output_dir, dataset, 'fnmr', f"{method_name}_{model}_{dataset}_fnmr4.npy"), fnmr4)
239 |
240 | save_pdf(fnmrs_list_2, method_labels, model=model, output_dir=args.output_dir, fmr =1e-2, db=dataset)
241 | save_pdf(fnmrs_list_3, method_labels, model=model, output_dir=args.output_dir, fmr =1e-3, db=dataset)
242 | save_pdf(fnmrs_list_4, method_labels, model=model, output_dir=args.output_dir, fmr =1e-4, db=dataset)
243 |
244 | def main():
245 | args = parser.parse_args()
246 | perform_1v1_quality_eval(args)
247 |
248 | if __name__ == '__main__':
249 | main()
250 |
--------------------------------------------------------------------------------
/ERC/erc_ijbc.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 |
3 | import argparse
4 | import os
5 | import glob
6 | import math
7 | import numpy as np
8 | from tqdm import tqdm
9 | import sklearn
10 | from sklearn import metrics
11 |
12 | from collections import defaultdict
13 | import gc
14 |
15 | import matplotlib.pyplot as plt
16 | from matplotlib.font_manager import FontProperties
17 |
18 | from roc import get_eer_threshold
19 |
20 | parser = argparse.ArgumentParser(description='Evaluation')
21 | parser.add_argument('--embeddings_dir', type=str,
22 | default="/usr/quality_embeddings",
23 | help='The dir save embeddings for each method and dataset, the diretory inside should be: {dataset}_{model}, e.g., IJBC_ArcFaceModel')
24 | parser.add_argument('--quality_score_dir', type=str,
25 | default="/usr/quality_score",
26 | help='The dir save file of quality scores for each dataset and method, the file inside should be: {method}_{dataset}.txt, e.g., CRFIQAS_IJBC.txt')
27 |
28 | parser.add_argument('--models', type=str,
29 | default="ArcFaceModel, ElasticFaceModel, MagFaceModel, CurricularFaceModel",
30 | help='The evaluated FR model')
31 | parser.add_argument('--eval_db', type=str,
32 | default="IJBC",
33 | help='The evaluated dataset')
34 |
35 | parser.add_argument('--output_dir', type=str,
36 | default="erc_plot_auc_test_all",
37 | help='')
38 | parser.add_argument('--method_name', type=str,
39 | default="Serfiq,BRISQUE,SDD,rankIQ,magface,FaceQnet,DeepIQA,PFE,rankIQA,CRFIQAL,CRFIQAS"
40 | help='The evaluated image quality estimation method')
41 |
42 | parser.add_argument('--distance_metric', type=str,
43 | default='cosine',
44 | help='Euclidian Distance or Cosine Distance.')
45 | parser.add_argument('--feat_size', type=int,
46 | default=1024,
47 | help='The size of extracted features')
48 |
49 | def load_all_features(root):
50 | all_features = defaultdict()
51 | for feature_path in tqdm(glob.glob(os.path.join(root, '*.npy'))):
52 | feat = np.load(feature_path)
53 | all_features[os.path.basename(feature_path)] = feat
54 | print("All features are loaded")
55 | return all_features
56 |
57 | def load_ijbc_pairs_features(pair_path, all_features, hq_pairs_idx, feature_size=1024):
58 | with open(pair_path, 'r') as f:
59 | lines=f.readlines()
60 |
61 | # build two empty embeddings matrix
62 | embeddings_0, embeddings_1 = np.empty([hq_pairs_idx.shape[0], feature_size]), np.empty([hq_pairs_idx.shape[0], feature_size])
63 | # load embeddings based on the needed pairs
64 | for indx in tqdm(range(hq_pairs_idx.shape[0])):
65 | real_index = hq_pairs_idx[indx]
66 | split_line = lines[real_index].split()
67 | feat_a = all_features[(split_line[0] + '.npy')]
68 | feat_b = all_features[(split_line[1] + '.npy')]
69 | embeddings_0[indx] = np.asarray(feat_a, dtype=np.float64)
70 | embeddings_1[indx] = np.asarray(feat_b, dtype=np.float64)
71 |
72 | return embeddings_0, embeddings_1
73 |
74 | def load_ijbc_pairs_quality(pair_path):
75 | with open(pair_path, 'r') as f:
76 | lines=f.readlines()
77 | pairs_quality, targets = [], []
78 |
79 | for idex, line in enumerate(tqdm(lines)):
80 | split_line = line.split()
81 | pairs_quality.append(float(split_line[3])) # quality score
82 | targets.append(int(split_line[2])) # imposter or genuine
83 | targets = np.vstack(targets).reshape(-1, )
84 | print('Loaded quality score and target for each pair')
85 | return targets, np.array(pairs_quality)
86 |
87 | def save_pdf(fnmrs_lists, method_labels, model, output_dir, fmr, db):
88 | fontsize = 20
89 | colors = ['green', 'black', 'orange', 'plum', 'cyan', 'gold', 'gray', 'salmon', 'deepskyblue', 'red', 'blue',
90 | 'darkseagreen', 'seashell', 'hotpink', 'indigo', 'lightseagreen', 'khaki', 'brown', 'teal', 'darkcyan']
91 | STYLES = ['--', '-.', ':', 'v--', '^--', ',--', '<--', '>--', '1--',
92 | '-' ,'-' , '2--', '3--', '4--', '.--', 'p--', '*--', 'h--', 'H--', '+--', 'x--', 'd--', '|--', '---']
93 |
94 | unconsidered_rates = 100 * np.arange(0, 0.98, 0.05)
95 | fig = plt.figure(figsize=(8, 6))
96 | ax1 = fig.add_subplot(111)
97 | if (not os.path.isdir(os.path.join(output_dir, 'plots', db))):
98 | os.makedirs(os.path.join(output_dir, 'plots', db))
99 |
100 | for i in range(len(fnmrs_lists)):
101 | ax1.plot(unconsidered_rates[:len(fnmrs_lists[i])], fnmrs_lists[i], STYLES[i], color=colors[i],
102 | label=method_labels[i], linewidth=4, markersize=12)
103 | auc = metrics.auc( np.array(unconsidered_rates[:len(fnmrs_lists[i])]/100),np.array(fnmrs_lists[i]))
104 | with open(os.path.join(output_dir, db, str(fmr)+"_auc.txt"), "a") as f:
105 | f.write(db + ':' + model + ':' + method_labels[i] + ':' + str(auc) + '\n')
106 |
107 | plt.xlabel('Ratio of unconsidered image [%]', fontsize=fontsize)
108 | plt.xlim([0, 98])
109 | plt.xticks(np.arange(0, 98, 10), fontsize=fontsize)
110 | plt.title(f"Testing on {db}, FMR={fmr}" + f" ({model})", fontsize=fontsize) # update : -3
111 | plt.ylabel('FNMR', fontsize=fontsize)
112 |
113 | plt.yticks(fontsize=fontsize)
114 | plt.grid(alpha=0.2)
115 |
116 | axbox = ax1.get_position() #
117 | plt.legend(bbox_to_anchor=(axbox.x0 + 0.5 * axbox.width, axbox.y0 - 0.22), prop=FontProperties(size=12),
118 | loc='lower center', ncol=6) #
119 | #ax1.get_legend().remove()
120 | plt.tight_layout()
121 | plt.savefig(os.path.join(output_dir, 'plots', db, db + '_' + str(fmr) +'_'+model + '.pdf'), bbox_inches='tight')
122 |
123 | def perform_1v1_quality_eval(args):
124 | d = ['IJBC']
125 | d = args.eval_db.split(',')
126 |
127 | for dataset in d:
128 | if os.path.exists(os.path.join(args.output_dir, dataset, str(1e-2)+"_auc.txt")):
129 | os.remove(os.path.join(args.output_dir, dataset, str(1e-2)+"_auc.txt"))
130 | if os.path.exists(os.path.join(args.output_dir, dataset, str(1e-3)+"_auc.txt")):
131 | os.remove(os.path.join(args.output_dir, dataset, str(1e-3)+"_auc.txt"))
132 | if os.path.exists(os.path.join(args.output_dir, dataset, str(1e-4)+"_auc.txt")):
133 | os.remove(os.path.join(args.output_dir, dataset, str(1e-4)+"_auc.txt"))
134 |
135 | match = True
136 | models=args.models.split(',')
137 | for model in models:
138 | for dataset in d:
139 | # create empty list for saving results at fnmr=1e-2, fnmr=1e-3, fnmr=1e-4,
140 | fnmrs_list_2, fnmrs_list_3, fnmrs_list_4, method_labels = [], [], []
141 | method_labels=[]
142 | method_names = args.method_name.split(',')
143 |
144 | if (not os.path.isdir(os.path.join(args.output_dir, dataset, 'fnmr'))):
145 | os.makedirs(os.path.join(args.output_dir, dataset, 'fnmr'))
146 |
147 | # 1. load all features based on number of images
148 | all_features = load_all_features(root=os.path.join(args.embeddings_dir, f"{dataset}_{model}"))
149 | unconsidered_rates = np.arange(0, 0.98, 0.05)
150 | desc = True
151 | for method_name in method_names:
152 | print(f"----process {model} {dataset} {method_name}-----------")
153 | targets, qlts = load_ijbc_pairs_quality(os.path.join(args.quality_score_dir, f"{method_name}_{dataset}.txt"))
154 |
155 | desc = False if method_name == 'PFE'
156 |
157 | if (desc):
158 | qlts_sorted_idx = np.argsort(qlts) # [::-1]
159 | else:
160 | qlts_sorted_idx = np.argsort(qlts)[::-1]
161 |
162 | num_pairs = len(targets)
163 | fnmrs_list_2_inner = []
164 | fnmrs_list_3_inner = []
165 | fnmrs_list_4_inner = []
166 |
167 | for u_rate in tqdm(unconsidered_rates):
168 | # compute the used paris based on unconsidered rates
169 | hq_pairs_idx = qlts_sorted_idx[int(u_rate * num_pairs):]
170 |
171 | # load features based on hq_pairs_idx
172 | x, y = load_ijbc_pairs_features(os.path.join(args.quality_score_dir, f"{method_name}_{dataset}.txt"), all_features, hq_pairs_idx, args.feat_size)
173 |
174 | print('Calculate distance....')
175 | if args.distance_metric == 'cosine':
176 | dot = np.sum(np.multiply(x, y), axis=1)
177 | norm = np.linalg.norm(x, axis=1) * np.linalg.norm(y, axis=1)
178 | similarity = np.clip(dot/norm, -1., 1.)
179 | dist = np.arccos(similarity) / math.pi
180 | del dot, norm, similarity, x, y
181 | gc.collect()
182 | else:
183 | x = sklearn.preprocessing.normalize(x)
184 | y = sklearn.preprocessing.normalize(y)
185 | diff = np.subtract(x, y)
186 | dist = np.sum(np.square(diff), 1)
187 | del diff, x, y
188 | gc.collect()
189 |
190 | # sort in a desending order
191 | pos_dists =np.sort(dist[targets[hq_pairs_idx] == 1])
192 | neg_dists = np.sort(dist[targets[hq_pairs_idx] == 0])
193 | print('Compute threshold......')
194 | fmr100_th, fmr1000_th, fmr10000_th = get_eer_threshold(pos_dists, neg_dists, ds_scores=True)
195 |
196 | g_true = [g for g in pos_dists if g < fmr100_th]
197 | fnmrs_list_2_inner.append(1- len(g_true)/(len(pos_dists)))
198 | g_true = [g for g in pos_dists if g < fmr1000_th]
199 | fnmrs_list_3_inner.append(1 - len(g_true) / (len(pos_dists)))
200 | g_true = [g for g in pos_dists if g < fmr10000_th]
201 | fnmrs_list_4_inner.append(1 - len(g_true) / (len(pos_dists)))
202 | del pos_dists, neg_dists
203 | gc.collect()
204 |
205 | np.save(os.path.join(args.output_dir, dataset, 'fnmr', f"{method_name}_{model}_{dataset}_fnmr2.npy"), fnmrs_list_2_inner)
206 | np.save(os.path.join(args.output_dir, dataset, 'fnmr', f"{method_name}_{model}_{dataset}_fnmr3.npy"), fnmrs_list_3_inner)
207 | np.save(os.path.join(args.output_dir, dataset, 'fnmr', f"{method_name}_{model}_{dataset}_fnmr4.npy"), fnmrs_list_4_inner)
208 |
209 | fnmrs_list_2.append(fnmrs_list_2_inner)
210 | fnmrs_list_3.append(fnmrs_list_3_inner)
211 | fnmrs_list_4.append(fnmrs_list_4_inner)
212 | method_labels.append(f"{method_name}")
213 |
214 | save_pdf(fnmrs_list_2, method_labels, model=model, output_dir=args.output_dir, fmr=1e-2, db=dataset)
215 | save_pdf(fnmrs_list_3, method_labels, model=model, output_dir=args.output_dir, fmr=1e-3, db=dataset)
216 | save_pdf(fnmrs_list_4, method_labels, model=model, output_dir=args.output_dir, fmr=1e-4, db=dataset)
217 |
218 | def main():
219 | args = parser.parse_args()
220 | perform_1v1_quality_eval(args)
221 |
222 | if __name__ == '__main__':
223 | main()
224 |
--------------------------------------------------------------------------------
/ERC/roc.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import operator
3 |
4 | def calculate_roc(gscores, iscores, ds_scores=False, rates=True):
5 | """Calculates FMR, FNMR
6 | @param gscores: Genuine matching scores
7 | @type gscores: Union[list, ndarray]
8 | @param iscores: Impostor matching scores
9 | @type giscores: Union[list, ndarray]
10 | @param ds_scores: Indicates whether input scores are
11 | dissimilarity scores
12 | @type ds_scores: bool
13 | @param rates: Indicates whether to return error rates instead
14 | of error values
15 | @type rates: bool
16 | @return: (thresholds, FMR, FNMR) or (thresholds, FM, FNM)
17 | @rtype: tuple
18 | """
19 | if isinstance(gscores, list):
20 | gscores = np.array(gscores, dtype=np.float64)
21 |
22 | if isinstance(iscores, list):
23 | iscores = np.array(iscores, dtype=np.float64)
24 |
25 | if gscores.dtype == np.int:
26 | gscores = np.float64(gscores)
27 |
28 | if iscores.dtype == np.int:
29 | iscores = np.float64(iscores)
30 |
31 | if ds_scores:
32 | gscores = gscores * -1
33 | iscores = iscores * -1
34 |
35 | gscores_number = len(gscores)
36 | iscores_number = len(iscores)
37 |
38 | # Labeling genuine scores as 1 and impostor scores as 0
39 | gscores = zip(gscores, [1] * gscores_number)
40 | iscores = zip(iscores, [0] * iscores_number)
41 |
42 | # Python3 compatibility
43 | gscores = list(gscores)
44 | iscores = list(iscores)
45 |
46 | # Stacking scores
47 | scores = np.array(sorted(gscores + iscores, key=operator.itemgetter(0)))
48 | cumul = np.cumsum(scores[:, 1])
49 |
50 | # Grouping scores
51 | thresholds, u_indices = np.unique(scores[:, 0], return_index=True)
52 |
53 | # Calculating FNM and FM distributions
54 | fnm = cumul[u_indices] - scores[u_indices][:, 1] # rejecting s < t
55 | fm = iscores_number - (u_indices - fnm)
56 |
57 | # Calculating FMR and FNMR
58 | if rates:
59 | fnm_rates = fnm / gscores_number
60 | fm_rates = fm / iscores_number
61 | else:
62 | fnm_rates = fnm
63 | fm_rates = fm
64 |
65 | if ds_scores:
66 | return thresholds * -1, fm_rates, fnm_rates
67 |
68 | return thresholds, fm_rates, fnm_rates
69 |
70 | def get_fmr_op(fmr, fnmr, op):
71 | """Returns the value of the given FMR operating point
72 | Definition:
73 | ZeroFMR: is defined as the lowest FNMR at which no false matches occur.
74 | Others FMR operating points are defined in a similar way.
75 | @param fmr: False Match Rates
76 | @type fmr: ndarray
77 | @param fnmr: False Non-Match Rates
78 | @type fnmr: ndarray
79 | @param op: Operating point
80 | @type op: float
81 | @returns: Index, The lowest FNMR at which the probability of FMR == op
82 | @rtype: float
83 | """
84 | index = np.argmin(abs(fmr - op))
85 | return index, fnmr[index]
86 |
87 | def get_fnmr_op(fmr, fnmr, op):
88 | """Returns the value of the given FNMR operating point
89 | Definition:
90 | ZeroFNMR: is defined as the lowest FMR at which no non-false matches occur.
91 | Others FNMR operating points are defined in a similar way.
92 | @param fmr: False Match Rates
93 | @type fmr: ndarray
94 | @param fnmr: False Non-Match Rates
95 | @type fnmr: ndarray
96 | @param op: Operating point
97 | @type op: float
98 | @returns: Index, The lowest FMR at which the probability of FNMR == op
99 | @rtype: float
100 | """
101 | temp = abs(fnmr - op)
102 | min_val = np.min(temp)
103 | index = np.where(temp == min_val)[0][-1]
104 | #index = np.argmin(abs(fnmr - op))
105 |
106 | return index, fmr[index]
107 |
108 | def get_eer_threshold(gen_scores, imp_scores, hformat=False, ds_scores=False):
109 | """Calculates EER associated statistics
110 | Keyword Arguments:
111 | @param gen_scores: The genuine scores
112 | @type gen_scores: list
113 | @param imp_scores: The impostor scores
114 | @type imp_scores: list
115 | @param id: An id for the experiment
116 | @type id: str
117 | @param hformat: Indicates whether the impostor scores are in histogram
118 | format
119 | @type hformat: bool
120 | @param ds_scores: Indicates whether the input scores are dissimilarity
121 | scores
122 | @type ds_scores: bool
123 | """
124 |
125 | # Calculating probabilities using scores as thrs
126 | roc_info = calculate_roc(gen_scores, imp_scores,
127 | ds_scores, rates=False)
128 | gnumber = len(gen_scores)
129 | inumber = len(imp_scores)
130 | thrs, fm, fnm = roc_info
131 | fmr = fm / inumber
132 | fnmr = fnm / gnumber
133 | ind, fmr1000 = get_fmr_op(fmr, fnmr, 0.001)
134 | fmr1000_th = thrs[ind]
135 |
136 | ind, fmr100 = get_fmr_op(fmr, fnmr, 0.01)
137 | fmr100_th = thrs[ind]
138 |
139 | ind, fmr10000 = get_fmr_op(fmr, fnmr, 0.0001)
140 | fmr10000_th = thrs[ind]
141 |
142 |
143 |
144 |
145 |
146 |
147 |
148 | return fmr100_th, fmr1000_th, fmr10000_th
149 |
--------------------------------------------------------------------------------
/IJB-C Evaluation/Readme.md:
--------------------------------------------------------------------------------
1 | # Evaluation
2 |
3 | ## Evaluation on IJB-C
4 | Follow this instruction to reproduce the results on IJB-C.
5 | 1. Download [IJB-C](https://www.nist.gov/programs-projects/face-challenges)
6 | 2. Unzip IJB-C (Folder structure should look like this /data/IJB_release/IJBC/)
7 | 3. Set (in extract_IJB.py) path = "/data/IJB_release/IJBC" to your IJB-C data folder and outpath = "/data/quality_data" where you want to save the preprocessed data
8 | 4. Run python extract_IJB.py (it creates the output folder, aligns and saves the images in BGR format, and creates image_path_list.txt and pair_list.txt)
9 | 5. Run xx to estimate the quality scores
10 | 5.1 For CR-FIQA(L)
11 | 5.2 For CR-FIQA(S)
12 | 6. Run CUDA_VISIBLE_DEVICES=0 python extract_emb.py --model_path ./pretrained/ElasticFace --model_id 295672 --dataset_path ./data/quality_data/IJBC
13 | #### Aggregate IJB-C Templates and Create new Pair List
14 | 7. Copy CR-FIQA(L) scores to CR-FIQA/feature_extraction/quality_scores/CRFIQAL_IJBC.txt
15 | 8. (rename /data/quality_embeddings/IJBC_ElasticFaceModel to /data/quality_embeddings/IJBC_ElasticFaceModel_raw if not already done by previous script)
16 | 9. Run python ijb_qs_pair_file.py --dataset_path /data/IJB_release/IJBC --q_modelname CRFIQAL (it aggregates templates to match the pair file, saves them at /data/quality_embeddings/IJBC_ElasticFaceModel/ and creates a pair file with quality scores for each aggregated template and saves it at /data/quality_data/IJBC)
17 | #### Evaluation with Quality Scores
18 | 10. Create folder CR-FIQA/feature_extraction/results/
19 | 11. CUDA_VISIBLE_DEVICES=0 python eval_ijbc_qs.py --model_path ./pretrained/ElasticFace --model_id 295672 --image-path /data/IJB_release/IJBC
20 |
--------------------------------------------------------------------------------
/IJB-C Evaluation/extract_IJB.py:
--------------------------------------------------------------------------------
1 | import os
2 | import cv2
3 | import numpy as np
4 | import shutil
5 | from skimage import transform as trans
6 | from tqdm import tqdm
7 |
8 |
9 | # change this for other dataset
10 | path = "/data/fboutros/IJB_release/IJB_release/IJBC"
11 | image_size = (112,112)
12 | outpath = "./data/quality_data"
13 |
14 | ref_lmk = np.array([
15 | [30.2946, 51.6963],
16 | [65.5318, 51.5014],
17 | [48.0252, 71.7366],
18 | [33.5493, 92.3655],
19 | [62.7299, 92.2041]], dtype=np.float32)
20 | ref_lmk[:, 0] += 8.0
21 |
22 | dataset_name = path.split("/")[-1]
23 | rel_img_path = os.path.join(outpath.split("/")[-1], dataset_name, "images")
24 | outpath = os.path.join(outpath, dataset_name)
25 | if not os.path.exists(outpath):
26 | os.makedirs(outpath)
27 | os.makedirs(os.path.join(outpath, "images"))
28 |
29 | print("extract:", dataset_name)
30 |
31 | img_path = os.path.join(path, "loose_crop")
32 | img_list_path = os.path.join(path, "meta", f"{dataset_name.lower()}_name_5pts_score.txt")
33 | img_list = open(img_list_path)
34 | files_list = img_list.readlines()
35 |
36 | txt_file = open(os.path.join(outpath, "image_path_list.txt"), "w")
37 |
38 | for img_index, each_line in tqdm(enumerate(files_list), total=len(files_list)):
39 | name_lmk_score = each_line.strip().split(' ')
40 | img_name = os.path.join(img_path, name_lmk_score[0])
41 | img = cv2.imread(img_name)
42 | lmk = np.array([float(x) for x in name_lmk_score[1:-1]],
43 | dtype=np.float32)
44 | lmk = lmk.reshape((5, 2))
45 |
46 | assert lmk.shape[0] == 5 and lmk.shape[1] == 2
47 |
48 | tform = trans.SimilarityTransform()
49 | tform.estimate(lmk, ref_lmk)
50 | M = tform.params[0:2, :]
51 | img = cv2.warpAffine(img, M, image_size, borderValue=0.0)
52 | img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
53 |
54 | cv2.imwrite(os.path.join(outpath, "images", name_lmk_score[0]), img)
55 | txt_file.write(os.path.join(rel_img_path, name_lmk_score[0])+"\n")
56 |
57 |
58 | txt_file.close()
59 | shutil.copy(
60 | os.path.join(path, "meta", f"{dataset_name.lower()}_template_pair_label.txt"),
61 | os.path.join(outpath, "pair_list.txt")
62 | )
63 | print("pair_list saved")
64 |
--------------------------------------------------------------------------------
/IJB-C Evaluation/extract_emb.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import os
3 | import sys
4 |
5 | import numpy as np
6 | from tqdm import tqdm
7 |
8 |
9 | def parse_arguments(argv):
10 | parser = argparse.ArgumentParser()
11 |
12 | parser.add_argument('--dataset_path', type=str, default='./data/quality_data/XQLFW',
13 | help='dataset directory')
14 | parser.add_argument('--modelname', type=str, default='ElasticFaceModel',
15 | help='ArcFaceModel, CurricularFaceModel, ElasticFaceModel, MagFaceModel')
16 | parser.add_argument('--gpu_id', type=int, default=0,
17 | help='GPU id.')
18 | parser.add_argument('--model_path', type=str, default="",
19 | help='path to pretrained model.')
20 | parser.add_argument('--model_id', type=str, default="295672",
21 | help='digit number in backbone file name')
22 | parser.add_argument('--relative_dir', type=str, default='./data',
23 | help='path to save the embeddings')
24 | return parser.parse_args(argv)
25 |
26 |
27 |
28 | def read_img_path_list(image_path_file, relative_dir):
29 | with open(image_path_file, "r") as f:
30 | lines = f.readlines()
31 | lines = [os.path.join(relative_dir, line.rstrip()) for line in lines]
32 | return lines
33 |
34 | def main(param):
35 | dataset_path = param.dataset_path
36 | modelname = param.modelname
37 | gpu_id = param.gpu_id
38 | data_path = param.relative_dir
39 | dataset_name = dataset_path.split("/")[-1]
40 | out_path = os.path.join(data_path, "quality_embeddings", f"{dataset_name}_{modelname}")
41 | if "ijb" in dataset_name.lower():
42 | out_path += "_raw"
43 | if not os.path.exists(out_path):
44 | os.makedirs(out_path)
45 | image_path_list = read_img_path_list(os.path.join(dataset_path, "image_path_list.txt"), data_path)
46 | if modelname == "ArcFaceModel":
47 | from model.ArcFaceModel import ArcFaceModel
48 | face_model = ArcFaceModel(param.model_path, param.model_id, gpu_id=gpu_id)
49 | elif modelname == "CurricularFaceModel":
50 | from model.CurricularFaceModel import CurricularFaceModel
51 | face_model = CurricularFaceModel(param.model_path, param.model_id, gpu_id)
52 | elif modelname == "ElasticFaceModel":
53 | from model.ElasticFaceModel import ElasticFaceModel
54 | face_model = ElasticFaceModel(param.model_path, param.model_id, gpu_id)
55 | elif modelname == "MagFaceModel":
56 | from model.MagFaceModel import MagFaceModel
57 | face_model = MagFaceModel(param.model_path, param.model_id, gpu_id)
58 | else:
59 | print("Unknown model")
60 | exit()
61 |
62 | features = face_model.get_batch_feature(image_path_list)
63 | features_flipped = face_model.get_batch_feature(image_path_list, flip=1)
64 |
65 | # too slow for IJBC
66 | # conc_features = np.concatenate((features, features_flipped), axis=1)
67 | # print(conc_features.shape)
68 |
69 | print("save features")
70 | for i in tqdm(range(len(features))):
71 | conc_features = np.concatenate((features[i], features_flipped[i]), axis=0)
72 | filename = str(str(image_path_list[i]).split("/")[-1].split(".")[0])
73 | np.save(os.path.join(out_path, filename), conc_features)
74 |
75 | if __name__ == '__main__':
76 | main(parse_arguments(sys.argv[1:]))
77 |
--------------------------------------------------------------------------------
/IJB-C Evaluation/ijb_qs_pair_file.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | import pandas as pd
4 | import numpy as np
5 | import sklearn.preprocessing
6 | from tqdm import tqdm
7 | import argparse
8 | from multiprocessing import Process, Queue
9 |
10 | def parse_arguments(argv):
11 | parser = argparse.ArgumentParser()
12 |
13 | parser.add_argument('--dataset_path', type=str, default='/data/fboutros/IJB_release/IJB_release/IJBC',
14 | help='dataset directory')
15 | parser.add_argument('--modelname', type=str, default='ElasticFaceModel',
16 | help='ArcFaceModel, CurricularFaceModel, ElasticFaceModel, MagFaceModel')
17 | parser.add_argument('--q_modelname', type=str, default='CRFIQAL',
18 | help='CRFIQAL, CRFIQAS')
19 | parser.add_argument('--threads', type=int, default=20,
20 | help='Number of threads')
21 | parser.add_argument('--raw_feature_path', type=str, default='./data/quality_embeddings',
22 | help='path to raw embeddings')
23 | parser.add_argument('--new_pair_path', type=str, default='./data/quality_data',
24 | help='path to save new quality pair list')
25 | return parser.parse_args(argv)
26 |
27 |
28 | def read_template_pair_list(path):
29 | pairs = pd.read_csv(path, sep=' ', header=None).values
30 | t1 = pairs[:, 0].astype(np.int)
31 | t2 = pairs[:, 1].astype(np.int)
32 | label = pairs[:, 2].astype(np.int)
33 | print("Template pair list loaded")
34 | return t1, t2, label
35 |
36 |
37 | def read_template_media_list(path):
38 | ijb_meta = pd.read_csv(path, sep=' ', header=None).values
39 | img_names = ijb_meta[:, 0]
40 | templates = ijb_meta[:, 1].astype(np.int)
41 | medias = ijb_meta[:, 2].astype(np.int)
42 | print("Tid mid list loaded")
43 | return img_names, templates, medias
44 |
45 |
46 | def load_raw_templates(files, raw_feature_path):
47 | features = np.zeros((len(files), 1024))
48 | for i, file in enumerate(files):
49 | filepath = os.path.join(raw_feature_path, file.replace(".jpg", ".npy"))
50 | features[i] = np.load(filepath)
51 | return features
52 |
53 |
54 | def aggregate_templates(tid_mid_path, feat_outpath, raw_feature_path):
55 | imgs_names, templates, medias = read_template_media_list(tid_mid_path)
56 |
57 | unique_templates = np.unique(templates)
58 |
59 | for uqt in tqdm(unique_templates, total=len(unique_templates)):
60 |
61 | (ind_t,) = np.where(templates == uqt)
62 | face_norm_feats = load_raw_templates(imgs_names[ind_t], raw_feature_path)
63 | face_medias = medias[ind_t]
64 | unique_medias, unique_media_counts = np.unique(face_medias,
65 | return_counts=True)
66 | media_norm_feats = []
67 | for u, ct in zip(unique_medias, unique_media_counts):
68 | (ind_m,) = np.where(face_medias == u)
69 | if ct == 1:
70 | media_norm_feats += [face_norm_feats[ind_m]]
71 | else: # image features from the same video will be aggregated into one feature
72 | media_norm_feats += [
73 | np.mean(face_norm_feats[ind_m], axis=0, keepdims=True)
74 | ]
75 | media_norm_feats = np.array(media_norm_feats, dtype=np.float32)
76 | aggregated_feature = sklearn.preprocessing.normalize(np.sum(media_norm_feats, axis=0))
77 | np.save(os.path.join(feat_outpath, f"{uqt}.npy"), aggregated_feature[0])
78 |
79 |
80 |
81 |
82 | def load_quality_scores(path):
83 | quality_scores = pd.read_csv(path, sep=' ', header=None).values
84 | score_dict = { qs[0].split("/")[-1] : float(qs[1]) for qs in quality_scores }
85 | print("Quality scores loaded")
86 | return score_dict
87 |
88 |
89 | def get_score_part(t1s, t2s, labels, quality_model, templates, imgs_names, score_dict, queue):
90 |
91 | def get_min_score(person):
92 | (ind_t1,) = np.where(templates == person)
93 | scores = []
94 | for img_name in imgs_names[ind_t1]:
95 | scores.append(score_dict[img_name])
96 | return max(scores) if quality_model == "PFE" else min(scores)
97 |
98 | q_pair_list = ""
99 | for t1, t2, label in tqdm(zip(t1s, t2s, labels), total=len(t1s)):
100 | min_score_t1 = get_min_score(t1)
101 | min_score_t2 = get_min_score(t2)
102 | if quality_model == "PFE":
103 | min_score = max([min_score_t1, min_score_t2])
104 | else:
105 | min_score = min([min_score_t1, min_score_t2])
106 |
107 | q_pair_list += f"{t1} {t2} {label} {min_score}\n"
108 |
109 | queue.put(q_pair_list)
110 |
111 |
112 | def create_quality_list(pair_path, tid_mid_path, quality_score_path, quality_model, dataset_name, threads, new_pair_path):
113 | imgs_names, templates, _ = read_template_media_list(tid_mid_path)
114 | score_dict = load_quality_scores(quality_score_path)
115 |
116 | t1s, t2s, labels = read_template_pair_list(pair_path)
117 | part_idx = len(t1s) // threads
118 | print(f"Quality estimation model: {quality_model}\n{threads+1} Threads")
119 |
120 | q = Queue()
121 | processes = []
122 |
123 | for idx in range(threads):
124 | t1_part = t1s[idx*part_idx:(idx+1)*part_idx]
125 | t2_part = t2s[idx*part_idx:(idx+1)*part_idx]
126 | label_part = labels[idx*part_idx:(idx+1)*part_idx]
127 | p = Process(target=get_score_part, args=(t1_part, t2_part, label_part, quality_model, templates, imgs_names, score_dict, q))
128 | processes.append(p)
129 | p.start()
130 |
131 | t1_part = t1s[threads*part_idx:]
132 | t2_part = t2s[threads*part_idx:]
133 | label_part = labels[threads*part_idx:]
134 | p = Process(target=get_score_part, args=(t1_part, t2_part, label_part, quality_model, templates, imgs_names, score_dict, q))
135 | processes.append(p)
136 | p.start()
137 |
138 |
139 | save_path = os.path.join(new_pair_path, dataset_name)
140 | pair_list = open(os.path.join(save_path, f"quality_pair_list_{quality_model}_{dataset_name}.txt"), "w")
141 |
142 | for p in processes:
143 | ret = q.get() # will block
144 | pair_list.write(ret)
145 | for p in processes:
146 | p.join()
147 |
148 | pair_list.close()
149 |
150 |
151 | def main(param):
152 | dataset_path = param.dataset_path
153 | dataset_name = dataset_path.split('/')[-1]
154 | model = param.modelname
155 | quality_model = param.q_modelname
156 | threads = param.threads
157 | new_pair_list_path = param.new_pair_path
158 | quality_score_path = os.path.join("quality_scores", f"{quality_model}_{dataset_name}.txt")
159 | raw_feature_path = os.path.join(param.raw_feature_path, dataset_name + "_" + model + "_raw")
160 | outpath = raw_feature_path[:-4]
161 | if not os.path.exists(outpath):
162 | os.makedirs(outpath)
163 |
164 | tid_mid_path = os.path.join(dataset_path, "meta", f"{dataset_name.lower()}_face_tid_mid.txt")
165 | pair_path = os.path.join(dataset_path, "meta", f"{dataset_name.lower()}_template_pair_label.txt")
166 |
167 | aggregate_templates(tid_mid_path, outpath, raw_feature_path)
168 |
169 | create_quality_list(pair_path, tid_mid_path, quality_score_path, quality_model, dataset_name, threads, new_pair_list_path)
170 |
171 |
172 | if __name__ == '__main__':
173 | main(parse_arguments(sys.argv[1:]))
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
2 | #### This is the official repository of the paper:
3 | ### CR-FIQA: Face Image Quality Assessment by Learning Sample Relative Classifiability
4 | ### Paper accepted at [CVPR 2023](https://cvpr2023.thecvf.com/)
5 | Arxiv: [CR-FIQA](https://arxiv.org/abs/2112.06592)
6 | ### News: 𝗿𝗮𝗻𝗸𝗲𝗱 𝟭𝘀𝘁/𝟮𝗻𝗱 at 𝗡𝗜𝗦𝗧 𝗙𝗮𝗰𝗲 𝗔𝗻𝗮𝗹𝘆𝘀𝗶𝘀 𝗧𝗲𝗰𝗵𝗻𝗼𝗹𝗼𝗴𝘆 𝗘𝘃𝗮𝗹𝘂𝗮𝘁𝗶𝗼𝗻 (𝗙𝗔𝗧𝗘) 𝗤𝘂𝗮𝗹𝗶𝘁𝘆
7 |
8 | https://pages.nist.gov/frvt/html/frvt_quality.html
9 |
10 | #### Update
11 | - New paper accepted at [CVPRW 2024](https://cvpr2023.thecvf.com/): GraFIQs: Face Image Quality Assessment Using Gradient Magnitudes
12 | https://github.com/jankolf/GraFIQs
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 | ## CR-FIQA training ##
22 | 1. In the paper, we employ MS1MV2 as the training dataset for CR-FIQA(L) which can be downloaded from InsightFace (MS1M-ArcFace in DataZoo)
23 | 1. Download MS1MV2 dataset from [insightface](https://github.com/deepinsight/insightface/tree/master/recognition/_datasets_) on strictly follow the licence distribution
24 | 3. We use CASIA-WebFace as the training dataset for CR-FIQA(S) which can be downloaded from InsightFace (CASIA in DataZoo)
25 | 1. Download CASIA dataset from [insightface](https://github.com/deepinsight/insightface/tree/master/recognition/_datasets_) on strictly follow the licence distribution
26 | 4. Unzip the dataset and place it in the data folder
27 | 5. Intall the requirement from requirement.txt
28 | 6. pip install -r requirements.txt
29 | 7. All code are trained and tested using PyTorch 1.7.1
30 | Details are under (Torch)[https://pytorch.org/get-started/locally/]
31 |
32 |
33 | ### CR-FIQA(L) ###
34 | Set the following in the config.py
35 | 1. config.output to output dir
36 | 2. config.network = "iresnet100"
37 | 3. config.dataset = "emoreIresNet"
38 | 4. Run ./run.sh
39 |
40 |
41 | ### CR-FIQA(S) ###
42 | Set the following in the config.py
43 | 1. config.output to output dir
44 | 2. config.network = "iresnet50"
45 | 3. config.dataset = "webface"
46 | 4. Run ./run.sh
47 |
48 | ## Pretrained model
49 |
50 | #### [CR-FIQA(L)](https://drive.google.com/drive/folders/1siy_3eQSBuIV6U6_9wgGtbZG2GMgVLMy?usp=sharing)
51 |
52 |
53 | #### [CR-FIQA(S)](https://drive.google.com/drive/folders/13bE4LP303XA_IzL1YOgG5eN0c8efHU9h?usp=sharing)
54 |
55 | ## Evaluation ##
56 | Follow these steps to reproduce the results on XQLFW:
57 | 1. Download the **aligned** [XQLFW](https://martlgap.github.io/xqlfw/pages/download.html) (please download xqlfw_aligned_112.zip)
58 | 2. Note: if you have an unaligned dataset, please do alignment first using [Alignment](https://github.com/fdbtrs/SFace-Privacy-friendly-and-Accurate-Face-Recognition-using-Synthetic-Data/blob/master/utils/MTCNN_alignment_fast.py)
59 | 2. Unzip XQLFW (Folder structure should look like this ./data/XQLFW/xqlfw_aligned_112/)
60 | 3. Download also xqlfw_pairs.txt to ./data/XQLFW/xqlfw_pairs.txt
61 | 4. Set (in feature_extraction/extract_xqlfw.py) path = "./data/XQLFW" to your XQLFW data folder and outpath = "./data/quality_data" where you want to save the preprocessed data
62 | 5. Run python extract_xqlfw.py (it creates the output folder, saves the images in BGR format, creates image_path_list.txt and pair_list.txt)
63 | 6. Run evaluation/getQualityScore.py to estimate the quality scores
64 | 1. CR-FIQA(L)
65 | 1. Download the pretrained model
66 | 2. run: python3 evaluation/getQualityScorce.py --data_dir "./data/quality_data" --datasets "XQLFW" --model_path "path_to_pretrained_CF_FIQAL_model" --backbone "iresnet100" --model_id "181952" --score_file_name "CRFIQAL.txt"
67 | 2. CR-FIQA(S)
68 | 1. Download the pretrained model
69 | 2. run: python3 evaluation/getQualityScorce.py --data_dir "./data/quality_data" --datasets "XQLFW" --model_path "path_to_pretrained_CF_FIQAL_model" --backbone "iresnet50" --model_id "32572" --score_file_name "CRFIQAS.txt"
70 | #### Note: Our model process an aligned image using ([Alignment](https://github.com/fdbtrs/SFace-Privacy-friendly-and-Accurate-Face-Recognition-using-Synthetic-Data/blob/master/utils/MTCNN_alignment_fast.py)). The images should be of RGB color space. If you are using OpenCV to load the images, make sure to convert from BGR to RGB as OpenCV converts the color channel when it is loaded with OpenCV read command. All images are processed to have pixel values between -1 and 1 as in [QualityModel.py](https://github.com/fdbtrs/CR-FIQA/blob/main/evaluation/QualityModel.py)
71 | The quality score of LFW, AgeDB-30, CFP-FP, CALFW, CPLFW can be produced by following these steps:
72 | 1. LFW, AgeDB-30, CFP-FP, CALFW, CPLFW are be included in the training dataset folder [insightface](https://github.com/deepinsight/insightface/tree/master/recognition/_datasets_)
73 | 2. Set (in extract_bin.py) path = "/data/faces_emore/lfw.bin" to your LFW bin file and outpath = "./data/quality_data" where you want to save the preprocessed data (subfolder will be created)
74 | 3. Run python extract_bin.py (it creates the output folder, saves the images in BGR format, creates image_path_list.txt and pair_list.txt)
75 | 4. Run evaluation/getQualityScore.py to estimate the quality scores
76 | 1. CR-FIQA(L)
77 | 1. Download the pretrained model
78 | 2. run: python3 evaluation/getQualityScorce.py --data_dir "./data/quality_data" --datasets "XQLFW" --model_path "path_to_pretrained_CF_FIQAL_model" --backbone "iresnet100" --model_id "181952" --score_file_name "CRFIQAL.txt" --color_channel "BGR"
79 | 2. CR-FIQA(S)
80 | 1. Download the pretrained model
81 | 2. run: python3 evaluation/getQualityScorce.py --data_dir "./data/quality_data" --datasets "XQLFW" --model_path "path_to_pretrained_CF_FIQAL_model" --backbone "iresnet50" --model_id "32572" --score_file_name "CRFIQAS.txt" --color_channel "BGR"
82 |
83 |
84 | ## Ploting ERC curves ##
85 | 1. Download pretrained model e.g. [ElasticFace-Arc](https://github.com/fdbtrs/ElasticFace), [MagFac](https://github.com/IrvingMeng/MagFace), [CurricularFace](https://github.com/HuangYG123/CurricularFace) or [ArcFace](https://github.com/deepinsight/insightface)
86 | 2. Run CUDA_VISIBLE_DEVICES=0 python feature_extraction/extract_emb.py --model_path ./pretrained/ElasticFace --model_id 295672 --dataset_path "./data/quality_data/XQLFW" --modelname "ElasticFaceModel"
87 | 2.1 Note: change the path to pretrained model and other arguments according to the evaluated model
88 | 3. Run python3 ERC/erc.py (details in ERC/README.md)
89 |
90 |
91 | ## Citation ##
92 | If you use any of the code provided in this repository or the models provided, please cite the following paper:
93 | ```
94 | @InProceedings{Boutros_2023_CVPR,
95 | author = {Boutros, Fadi and Fang, Meiling and Klemt, Marcel and Fu, Biying and Damer, Naser},
96 | title = {CR-FIQA: Face Image Quality Assessment by Learning Sample Relative Classifiability},
97 | booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
98 | month = {June},
99 | year = {2023},
100 | pages = {5836-5845}
101 | }
102 | ```
103 |
104 |
105 | ## License ##
106 |
107 | This project is licensed under the terms of the Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) license.
108 | Copyright (c) 2021 Fraunhofer Institute for Computer Graphics Research IGD Darmstadt
109 |
--------------------------------------------------------------------------------
/backbones/__init__.py:
--------------------------------------------------------------------------------
1 |
2 |
--------------------------------------------------------------------------------
/backbones/iresnet.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch import nn
3 |
4 | __all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100']
5 |
6 |
7 | def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
8 | """3x3 convolution with padding"""
9 | return nn.Conv2d(in_planes,
10 | out_planes,
11 | kernel_size=3,
12 | stride=stride,
13 | padding=dilation,
14 | groups=groups,
15 | bias=False,
16 | dilation=dilation)
17 |
18 |
19 | def conv1x1(in_planes, out_planes, stride=1):
20 | """1x1 convolution"""
21 | return nn.Conv2d(in_planes,
22 | out_planes,
23 | kernel_size=1,
24 | stride=stride,
25 | bias=False)
26 | class SEModule(nn.Module):
27 | def __init__(self, channels, reduction):
28 | super(SEModule, self).__init__()
29 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
30 | self.fc1 = nn.Conv2d(channels, channels // reduction, kernel_size=1, padding=0, bias=False)
31 | self.relu = nn.ReLU(inplace=True)
32 | self.fc2 = nn.Conv2d(channels // reduction, channels, kernel_size=1, padding=0, bias=False)
33 | self.sigmoid = nn.Sigmoid()
34 |
35 | def forward(self, x):
36 | input = x
37 | x = self.avg_pool(x)
38 | x = self.fc1(x)
39 | x = self.relu(x)
40 | x = self.fc2(x)
41 | x = self.sigmoid(x)
42 |
43 | return input * x
44 |
45 | class IBasicBlock(nn.Module):
46 | expansion = 1
47 | def __init__(self, inplanes, planes, stride=1, downsample=None,
48 | groups=1, base_width=64, dilation=1,use_se=False):
49 | super(IBasicBlock, self).__init__()
50 | if groups != 1 or base_width != 64:
51 | raise ValueError('BasicBlock only supports groups=1 and base_width=64')
52 | if dilation > 1:
53 | raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
54 | self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,)
55 | self.conv1 = conv3x3(inplanes, planes)
56 | self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,)
57 | self.prelu = nn.PReLU(planes)
58 | self.conv2 = conv3x3(planes, planes, stride)
59 | self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,)
60 | self.downsample = downsample
61 | self.stride = stride
62 | self.use_se=use_se
63 | if (use_se):
64 | self.se_block=SEModule(planes,16)
65 |
66 | def forward(self, x):
67 | identity = x
68 | out = self.bn1(x)
69 | out = self.conv1(out)
70 | out = self.bn2(out)
71 | out = self.prelu(out)
72 | out = self.conv2(out)
73 | out = self.bn3(out)
74 | if(self.use_se):
75 | out=self.se_block(out)
76 | if self.downsample is not None:
77 | identity = self.downsample(x)
78 | out += identity
79 | return out
80 |
81 |
82 | class IResNet(nn.Module):
83 | fc_scale = 7 * 7
84 | def __init__(self,
85 | block, layers, dropout=0, num_features=512, zero_init_residual=False,
86 | groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False,use_se=False, qs=1):
87 | super(IResNet, self).__init__()
88 | self.fp16 = fp16
89 | self.inplanes = 64
90 | self.dilation = 1
91 | self.use_se=use_se
92 | self.qs=qs
93 | if replace_stride_with_dilation is None:
94 | replace_stride_with_dilation = [False, False, False]
95 | if len(replace_stride_with_dilation) != 3:
96 | raise ValueError("replace_stride_with_dilation should be None "
97 | "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
98 | self.groups = groups
99 | self.base_width = width_per_group
100 | self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
101 | self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
102 | self.prelu = nn.PReLU(self.inplanes)
103 | self.layer1 = self._make_layer(block, 64, layers[0], stride=2 ,use_se=self.use_se)
104 | self.layer2 = self._make_layer(block,
105 | 128,
106 | layers[1],
107 | stride=2,
108 | dilate=replace_stride_with_dilation[0],use_se=self.use_se)
109 | self.layer3 = self._make_layer(block,
110 | 256,
111 | layers[2],
112 | stride=2,
113 | dilate=replace_stride_with_dilation[1] ,use_se=self.use_se)
114 | self.layer4 = self._make_layer(block,
115 | 512,
116 | layers[3],
117 | stride=2,
118 | dilate=replace_stride_with_dilation[2] ,use_se=self.use_se)
119 | self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,)
120 | self.dropout = nn.Dropout(p=dropout, inplace=True)
121 | self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)
122 | self.features = nn.BatchNorm1d(num_features, eps=1e-05)
123 | self.qs=nn.Linear(num_features,self.qs)
124 |
125 | for m in self.modules():
126 | if isinstance(m, nn.Conv2d):
127 | nn.init.normal_(m.weight, 0, 0.1)
128 | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
129 | nn.init.constant_(m.weight, 1)
130 | nn.init.constant_(m.bias, 0)
131 |
132 | if zero_init_residual:
133 | for m in self.modules():
134 | if isinstance(m, IBasicBlock):
135 | nn.init.constant_(m.bn2.weight, 0)
136 |
137 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False,use_se=False):
138 | downsample = None
139 | previous_dilation = self.dilation
140 | if dilate:
141 | self.dilation *= stride
142 | stride = 1
143 | if stride != 1 or self.inplanes != planes * block.expansion:
144 | downsample = nn.Sequential(
145 | conv1x1(self.inplanes, planes * block.expansion, stride),
146 | nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),
147 | )
148 | layers = []
149 | layers.append(
150 | block(self.inplanes, planes, stride, downsample, self.groups,
151 | self.base_width, previous_dilation,use_se=use_se))
152 | self.inplanes = planes * block.expansion
153 | for _ in range(1, blocks):
154 | layers.append(
155 | block(self.inplanes,
156 | planes,
157 | groups=self.groups,
158 | base_width=self.base_width,
159 | dilation=self.dilation,use_se=use_se))
160 |
161 | return nn.Sequential(*layers)
162 |
163 | def forward(self, x):
164 | with torch.cuda.amp.autocast(self.fp16):
165 | x = self.conv1(x)
166 | x = self.bn1(x)
167 | x = self.prelu(x)
168 | x = self.layer1(x)
169 | x = self.layer2(x)
170 | x = self.layer3(x)
171 | x = self.layer4(x)
172 | x = self.bn2(x)
173 | x = torch.flatten(x, 1)
174 | x = self.dropout(x)
175 | x = self.fc(x.float() if self.fp16 else x)
176 |
177 | x = self.features(x)
178 | qs = self.qs(x)
179 | return x, qs
180 |
181 |
182 | class IdentityIResNet(nn.Module):
183 | fc_scale = 7 * 7
184 | def __init__(self,
185 | block, layers, dropout=0, num_features=512, zero_init_residual=False,
186 | groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False,use_se=False):
187 | super(IdentityIResNet, self).__init__()
188 | self.fp16 = fp16
189 | self.inplanes = 64
190 | self.dilation = 1
191 | self.use_se=use_se
192 | if replace_stride_with_dilation is None:
193 | replace_stride_with_dilation = [False, False, False]
194 | if len(replace_stride_with_dilation) != 3:
195 | raise ValueError("replace_stride_with_dilation should be None "
196 | "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
197 | self.groups = groups
198 | self.base_width = width_per_group
199 | self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
200 | self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
201 | self.prelu = nn.PReLU(self.inplanes)
202 | self.layer1 = self._make_layer(block, 64, layers[0], stride=2 ,use_se=self.use_se)
203 | self.layer2 = self._make_layer(block,
204 | 128,
205 | layers[1],
206 | stride=2,
207 | dilate=replace_stride_with_dilation[0],use_se=self.use_se)
208 | self.layer3 = self._make_layer(block,
209 | 256,
210 | layers[2],
211 | stride=2,
212 | dilate=replace_stride_with_dilation[1] ,use_se=self.use_se)
213 | self.layer4 = self._make_layer(block,
214 | 512,
215 | layers[3],
216 | stride=2,
217 | dilate=replace_stride_with_dilation[2] ,use_se=self.use_se)
218 | self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,)
219 | self.dropout = nn.Dropout(p=dropout, inplace=True)
220 | self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)
221 | self.features = nn.BatchNorm1d(num_features, eps=1e-05)
222 | nn.init.constant_(self.features.weight, 1.0)
223 | self.features.weight.requires_grad = False
224 |
225 |
226 | for m in self.modules():
227 | if isinstance(m, nn.Conv2d):
228 | nn.init.normal_(m.weight, 0, 0.1)
229 | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
230 | nn.init.constant_(m.weight, 1)
231 | nn.init.constant_(m.bias, 0)
232 |
233 | if zero_init_residual:
234 | for m in self.modules():
235 | if isinstance(m, IBasicBlock):
236 | nn.init.constant_(m.bn2.weight, 0)
237 |
238 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False,use_se=False):
239 | downsample = None
240 | previous_dilation = self.dilation
241 | if dilate:
242 | self.dilation *= stride
243 | stride = 1
244 | if stride != 1 or self.inplanes != planes * block.expansion:
245 | downsample = nn.Sequential(
246 | conv1x1(self.inplanes, planes * block.expansion, stride),
247 | nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),
248 | )
249 | layers = []
250 | layers.append(
251 | block(self.inplanes, planes, stride, downsample, self.groups,
252 | self.base_width, previous_dilation,use_se=use_se))
253 | self.inplanes = planes * block.expansion
254 | for _ in range(1, blocks):
255 | layers.append(
256 | block(self.inplanes,
257 | planes,
258 | groups=self.groups,
259 | base_width=self.base_width,
260 | dilation=self.dilation,use_se=use_se))
261 |
262 | return nn.Sequential(*layers)
263 |
264 | def forward(self, x):
265 | with torch.cuda.amp.autocast(self.fp16):
266 | x = self.conv1(x)
267 | x = self.bn1(x)
268 | x = self.prelu(x)
269 | x = self.layer1(x)
270 | x = self.layer2(x)
271 | x = self.layer3(x)
272 | x = self.layer4(x)
273 | x = self.bn2(x)
274 | x = torch.flatten(x, 1)
275 | x = self.dropout(x)
276 | x = self.fc(x.float() if self.fp16 else x)
277 |
278 | x = self.features(x)
279 | return x
280 |
281 | class OnTopQS(nn.Module):
282 | def __init__(self,
283 | num_features=512):
284 | super(OnTopQS, self).__init__()
285 | self.qs=nn.Linear(num_features,1)
286 |
287 | def forward(self, x):
288 | return self.qs(x)
289 |
290 |
291 | def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
292 | model = IResNet(block, layers, **kwargs)
293 | if pretrained:
294 | raise ValueError()
295 | return model
296 |
297 | def _iresne_identity(arch, block, layers, pretrained, progress, **kwargs):
298 | model = IdentityIResNet(block, layers, **kwargs)
299 | if pretrained:
300 | raise ValueError()
301 | return model
302 |
303 |
304 |
305 |
306 | def iresnet18(pretrained=False, progress=True, **kwargs):
307 | return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained,
308 | progress, **kwargs)
309 |
310 |
311 | def iresnet34(pretrained=False, progress=True, **kwargs):
312 | return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained,
313 | progress, **kwargs)
314 |
315 |
316 | def iresnet50(pretrained=False, progress=True, **kwargs):
317 | return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,
318 | progress, **kwargs)
319 |
320 | def iresnet50_identity(pretrained=False, progress=True, **kwargs):
321 | return _iresne_identity('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,
322 | progress, **kwargs)
323 |
324 | def iresnet100(pretrained=False, progress=True, **kwargs):
325 | return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained,
326 | progress, **kwargs)
327 | def _test():
328 | import torch
329 |
330 | pretrained = False
331 |
332 | models = [
333 | iresnet100
334 | ]
335 |
336 | for model in models:
337 |
338 | net = model()
339 | print(net)
340 |
341 | if __name__ == "__main__":
342 | _test()
343 |
--------------------------------------------------------------------------------
/config.py:
--------------------------------------------------------------------------------
1 | from easydict import EasyDict as edict
2 |
3 | config = edict()
4 | config.dataset = "webface" # training dataset
5 | config.embedding_size = 512 # embedding size of evaluation
6 | config.momentum = 0.9
7 | config.weight_decay = 5e-4
8 | config.batch_size = 128 # batch size per GPU
9 | config.lr = 0.1
10 | config.output = "output/R50_CRFIQA" # train evaluation output folder
11 | config.global_step=0 # step to resume
12 | config.s=64.0
13 | config.m=0.50
14 | config.beta=0.5
15 |
16 |
17 |
18 | # type of network to train [ iresnet100 | iresnet50 ]
19 | config.network = "iresnet50"
20 |
21 |
22 |
23 |
24 | if config.dataset == "emoreIresNet":
25 | config.rec = "datafaces_emore"
26 | config.num_classes = 85742
27 | config.num_image = 5822653
28 | config.num_epoch = 18
29 | config.warmup_epoch = -1
30 | config.val_targets = ["lfw", "cfp_fp", "cfp_ff", "agedb_30", "calfw", "cplfw"]
31 | config.eval_step=5686
32 | def lr_step_func(epoch):
33 | return ((epoch + 1) / (4 + 1)) ** 2 if epoch < -1 else 0.1 ** len(
34 | [m for m in [8, 14,20,25] if m - 1 <= epoch]) # [m for m in [8, 14,20,25] if m - 1 <= epoch])
35 | config.lr_func = lr_step_func
36 |
37 | elif config.dataset == "webface":
38 | config.rec = "data/faces_webface_112x112"
39 | config.num_classes = 10572
40 | config.num_image = 501195
41 | config.num_epoch = 34 # [22, 30, 35] [22, 30, 40]
42 | config.warmup_epoch = -1
43 | config.val_targets = ["lfw", "cfp_fp", "agedb_30"]
44 | config.eval_step= 958 #33350
45 |
46 | def lr_step_func(epoch):
47 | return ((epoch + 1) / (4 + 1)) ** 2 if epoch < config.warmup_epoch else 0.1 ** len(
48 | [m for m in [20, 28, 32] if m - 1 <= epoch])
49 | config.lr_func = lr_step_func
50 |
--------------------------------------------------------------------------------
/dataset.py:
--------------------------------------------------------------------------------
1 | import numbers
2 | import os
3 | import queue as Queue
4 | import threading
5 |
6 | import mxnet as mx
7 | import numpy as np
8 | import torch
9 | from torch.utils.data import DataLoader, Dataset
10 | from torchvision import transforms
11 |
12 |
13 | class BackgroundGenerator(threading.Thread):
14 | def __init__(self, generator, local_rank, max_prefetch=6):
15 | super(BackgroundGenerator, self).__init__()
16 | self.queue = Queue.Queue(max_prefetch)
17 | self.generator = generator
18 | self.local_rank = local_rank
19 | self.daemon = True
20 | self.start()
21 |
22 | def run(self):
23 | torch.cuda.set_device(self.local_rank)
24 | for item in self.generator:
25 | self.queue.put(item)
26 | self.queue.put(None)
27 |
28 | def next(self):
29 | next_item = self.queue.get()
30 | if next_item is None:
31 | raise StopIteration
32 | return next_item
33 |
34 | def __next__(self):
35 | return self.next()
36 |
37 | def __iter__(self):
38 | return self
39 |
40 |
41 | class DataLoaderX(DataLoader):
42 | def __init__(self, local_rank, **kwargs):
43 | super(DataLoaderX, self).__init__(**kwargs)
44 | self.stream = torch.cuda.Stream(local_rank)
45 | self.local_rank = local_rank
46 |
47 | def __iter__(self):
48 | self.iter = super(DataLoaderX, self).__iter__()
49 | self.iter = BackgroundGenerator(self.iter, self.local_rank)
50 | self.preload()
51 | return self
52 |
53 | def preload(self):
54 | self.batch = next(self.iter, None)
55 | if self.batch is None:
56 | return None
57 | with torch.cuda.stream(self.stream):
58 | for k in range(len(self.batch)):
59 | self.batch[k] = self.batch[k].to(device=self.local_rank,
60 | non_blocking=True)
61 |
62 | def __next__(self):
63 | torch.cuda.current_stream().wait_stream(self.stream)
64 | batch = self.batch
65 | if batch is None:
66 | raise StopIteration
67 | self.preload()
68 | return batch
69 |
70 |
71 | class MXFaceDataset(Dataset):
72 | def __init__(self, root_dir, local_rank):
73 | super(MXFaceDataset, self).__init__()
74 | self.transform = transforms.Compose(
75 | [transforms.ToPILImage(),
76 | transforms.RandomHorizontalFlip(),
77 | transforms.ToTensor(),
78 | transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
79 | ])
80 | self.root_dir = root_dir
81 | self.local_rank = local_rank
82 | path_imgrec = os.path.join(root_dir, 'train.rec')
83 | path_imgidx = os.path.join(root_dir, 'train.idx')
84 | self.imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r')
85 | s = self.imgrec.read_idx(0)
86 | header, _ = mx.recordio.unpack(s)
87 | if header.flag > 0:
88 | self.header0 = (int(header.label[0]), int(header.label[1]))
89 | self.imgidx = np.array(range(1, int(header.label[0])))
90 | else:
91 | self.imgidx = np.array(list(self.imgrec.keys))
92 |
93 | def __getitem__(self, index):
94 | idx = self.imgidx[index]
95 | s = self.imgrec.read_idx(idx)
96 | header, img = mx.recordio.unpack(s)
97 | label = header.label
98 | if not isinstance(label, numbers.Number):
99 | label = label[0]
100 | label = torch.tensor(label, dtype=torch.long)
101 | sample = mx.image.imdecode(img).asnumpy()
102 | if self.transform is not None:
103 | sample = self.transform(sample)
104 | return sample, label
105 |
106 | def __len__(self):
107 | return len(self.imgidx)
108 |
--------------------------------------------------------------------------------
/eval/verification.py:
--------------------------------------------------------------------------------
1 | """Helper for evaluation on the Labeled Faces in the Wild dataset
2 | """
3 |
4 | # MIT License
5 | #
6 | # Copyright (c) 2016 David Sandberg
7 | #
8 | # Permission is hereby granted, free of charge, to any person obtaining a copy
9 | # of this software and associated documentation files (the "Software"), to deal
10 | # in the Software without restriction, including without limitation the rights
11 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12 | # copies of the Software, and to permit persons to whom the Software is
13 | # furnished to do so, subject to the following conditions:
14 | #
15 | # The above copyright notice and this permission notice shall be included in all
16 | # copies or substantial portions of the Software.
17 | #
18 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
24 | # SOFTWARE.
25 |
26 |
27 | import datetime
28 | import os
29 | import pickle
30 |
31 | import mxnet as mx
32 | import numpy as np
33 | import sklearn
34 | import torch
35 | from mxnet import ndarray as nd
36 | from scipy import interpolate
37 | from sklearn.decomposition import PCA
38 | from sklearn.model_selection import KFold
39 |
40 |
41 | class LFold:
42 | def __init__(self, n_splits=2, shuffle=False):
43 | self.n_splits = n_splits
44 | if self.n_splits > 1:
45 | self.k_fold = KFold(n_splits=n_splits, shuffle=shuffle)
46 |
47 | def split(self, indices):
48 | if self.n_splits > 1:
49 | return self.k_fold.split(indices)
50 | else:
51 | return [(indices, indices)]
52 |
53 |
54 | def calculate_roc(thresholds,
55 | embeddings1,
56 | embeddings2,
57 | actual_issame,
58 | nrof_folds=10,
59 | pca=0):
60 | assert (embeddings1.shape[0] == embeddings2.shape[0])
61 | assert (embeddings1.shape[1] == embeddings2.shape[1])
62 | nrof_pairs = min(len(actual_issame), embeddings1.shape[0])
63 | nrof_thresholds = len(thresholds)
64 | k_fold = LFold(n_splits=nrof_folds, shuffle=False)
65 |
66 | tprs = np.zeros((nrof_folds, nrof_thresholds))
67 | fprs = np.zeros((nrof_folds, nrof_thresholds))
68 | accuracy = np.zeros((nrof_folds))
69 | indices = np.arange(nrof_pairs)
70 |
71 | if pca == 0:
72 | diff = np.subtract(embeddings1, embeddings2)
73 | dist = np.sum(np.square(diff), 1)
74 |
75 | for fold_idx, (train_set, test_set) in enumerate(k_fold.split(indices)):
76 | if pca > 0:
77 | print('doing pca on', fold_idx)
78 | embed1_train = embeddings1[train_set]
79 | embed2_train = embeddings2[train_set]
80 | _embed_train = np.concatenate((embed1_train, embed2_train), axis=0)
81 | pca_model = PCA(n_components=pca)
82 | pca_model.fit(_embed_train)
83 | embed1 = pca_model.transform(embeddings1)
84 | embed2 = pca_model.transform(embeddings2)
85 | embed1 = sklearn.preprocessing.normalize(embed1)
86 | embed2 = sklearn.preprocessing.normalize(embed2)
87 | diff = np.subtract(embed1, embed2)
88 | dist = np.sum(np.square(diff), 1)
89 |
90 | # Find the best threshold for the fold
91 | acc_train = np.zeros((nrof_thresholds))
92 | for threshold_idx, threshold in enumerate(thresholds):
93 | _, _, acc_train[threshold_idx] = calculate_accuracy(
94 | threshold, dist[train_set], actual_issame[train_set])
95 | best_threshold_index = np.argmax(acc_train)
96 | for threshold_idx, threshold in enumerate(thresholds):
97 | tprs[fold_idx, threshold_idx], fprs[fold_idx, threshold_idx], _ = calculate_accuracy(
98 | threshold, dist[test_set],
99 | actual_issame[test_set])
100 | _, _, accuracy[fold_idx] = calculate_accuracy(
101 | thresholds[best_threshold_index], dist[test_set],
102 | actual_issame[test_set])
103 |
104 | tpr = np.mean(tprs, 0)
105 | fpr = np.mean(fprs, 0)
106 | return tpr, fpr, accuracy
107 |
108 |
109 | def calculate_accuracy(threshold, dist, actual_issame):
110 | predict_issame = np.less(dist, threshold)
111 | tp = np.sum(np.logical_and(predict_issame, actual_issame))
112 | fp = np.sum(np.logical_and(predict_issame, np.logical_not(actual_issame)))
113 | tn = np.sum(
114 | np.logical_and(np.logical_not(predict_issame),
115 | np.logical_not(actual_issame)))
116 | fn = np.sum(np.logical_and(np.logical_not(predict_issame), actual_issame))
117 |
118 | tpr = 0 if (tp + fn == 0) else float(tp) / float(tp + fn)
119 | fpr = 0 if (fp + tn == 0) else float(fp) / float(fp + tn)
120 | acc = float(tp + tn) / dist.size
121 | return tpr, fpr, acc
122 |
123 |
124 | def calculate_val(thresholds,
125 | embeddings1,
126 | embeddings2,
127 | actual_issame,
128 | far_target,
129 | nrof_folds=10):
130 | assert (embeddings1.shape[0] == embeddings2.shape[0])
131 | assert (embeddings1.shape[1] == embeddings2.shape[1])
132 | nrof_pairs = min(len(actual_issame), embeddings1.shape[0])
133 | nrof_thresholds = len(thresholds)
134 | k_fold = LFold(n_splits=nrof_folds, shuffle=False)
135 |
136 | val = np.zeros(nrof_folds)
137 | far = np.zeros(nrof_folds)
138 |
139 | diff = np.subtract(embeddings1, embeddings2)
140 | dist = np.sum(np.square(diff), 1)
141 | indices = np.arange(nrof_pairs)
142 |
143 | for fold_idx, (train_set, test_set) in enumerate(k_fold.split(indices)):
144 |
145 | # Find the threshold that gives FAR = far_target
146 | far_train = np.zeros(nrof_thresholds)
147 | for threshold_idx, threshold in enumerate(thresholds):
148 | _, far_train[threshold_idx] = calculate_val_far(
149 | threshold, dist[train_set], actual_issame[train_set])
150 | if np.max(far_train) >= far_target:
151 | f = interpolate.interp1d(far_train, thresholds, kind='slinear')
152 | threshold = f(far_target)
153 | else:
154 | threshold = 0.0
155 |
156 | val[fold_idx], far[fold_idx] = calculate_val_far(
157 | threshold, dist[test_set], actual_issame[test_set])
158 |
159 | val_mean = np.mean(val)
160 | far_mean = np.mean(far)
161 | val_std = np.std(val)
162 | return val_mean, val_std, far_mean
163 |
164 |
165 | def calculate_val_far(threshold, dist, actual_issame):
166 | predict_issame = np.less(dist, threshold)
167 | true_accept = np.sum(np.logical_and(predict_issame, actual_issame))
168 | false_accept = np.sum(
169 | np.logical_and(predict_issame, np.logical_not(actual_issame)))
170 | n_same = np.sum(actual_issame)
171 | n_diff = np.sum(np.logical_not(actual_issame))
172 | # print(true_accept, false_accept)
173 | # print(n_same, n_diff)
174 | val = float(true_accept) / float(n_same)
175 | far = float(false_accept) / float(n_diff)
176 | return val, far
177 |
178 |
179 | def evaluate(embeddings, actual_issame, nrof_folds=10, pca=0):
180 | # Calculate evaluation metrics
181 | thresholds = np.arange(0, 4, 0.01)
182 | embeddings1 = embeddings[0::2]
183 | embeddings2 = embeddings[1::2]
184 | tpr, fpr, accuracy = calculate_roc(thresholds,
185 | embeddings1,
186 | embeddings2,
187 | np.asarray(actual_issame),
188 | nrof_folds=nrof_folds,
189 | pca=pca)
190 | thresholds = np.arange(0, 4, 0.001)
191 | val, val_std, far = calculate_val(thresholds,
192 | embeddings1,
193 | embeddings2,
194 | np.asarray(actual_issame),
195 | 1e-3,
196 | nrof_folds=nrof_folds)
197 | return tpr, fpr, accuracy, val, val_std, far
198 |
199 | @torch.no_grad()
200 | def load_bin(path, image_size):
201 | try:
202 | with open(path, 'rb') as f:
203 | bins, issame_list = pickle.load(f) # py2
204 | except UnicodeDecodeError as e:
205 | with open(path, 'rb') as f:
206 | bins, issame_list = pickle.load(f, encoding='bytes') # py3
207 | data_list = []
208 | for flip in [0, 1]:
209 | data = torch.empty((len(issame_list) * 2, 3, image_size[0], image_size[1]))
210 | data_list.append(data)
211 | for idx in range(len(issame_list) * 2):
212 | _bin = bins[idx]
213 | img = mx.image.imdecode(_bin)
214 | if img.shape[1] != image_size[0]:
215 | img = mx.image.resize_short(img, image_size[0])
216 | img = nd.transpose(img, axes=(2, 0, 1))
217 | for flip in [0, 1]:
218 | if flip == 1:
219 | img = mx.ndarray.flip(data=img, axis=2)
220 | data_list[flip][idx][:] = torch.from_numpy(img.asnumpy())
221 | if idx % 1000 == 0:
222 | print('loading bin', idx)
223 | print(data_list[0].shape)
224 | return data_list, issame_list
225 |
226 | @torch.no_grad()
227 | def test(data_set, backbone, batch_size, nfolds=10):
228 | print('testing verification..')
229 | data_list = data_set[0]
230 | issame_list = data_set[1]
231 | embeddings_list = []
232 | time_consumed = 0.0
233 | for i in range(len(data_list)):
234 | data = data_list[i]
235 | embeddings = None
236 | ba = 0
237 | while ba < data.shape[0]:
238 | bb = min(ba + batch_size, data.shape[0])
239 | count = bb - ba
240 | _data = data[bb - batch_size: bb]
241 | time0 = datetime.datetime.now()
242 | img = ((_data / 255) - 0.5) / 0.5
243 | net_out ,qs_ = backbone(img)
244 | _embeddings = net_out.detach().cpu().numpy()
245 | time_now = datetime.datetime.now()
246 | diff = time_now - time0
247 | time_consumed += diff.total_seconds()
248 | if embeddings is None:
249 | embeddings = np.zeros((data.shape[0], _embeddings.shape[1]))
250 | embeddings[ba:bb, :] = _embeddings[(batch_size - count):, :]
251 | ba = bb
252 | embeddings_list.append(embeddings)
253 |
254 | _xnorm = 0.0
255 | _xnorm_cnt = 0
256 | for embed in embeddings_list:
257 | for i in range(embed.shape[0]):
258 | _em = embed[i]
259 | _norm = np.linalg.norm(_em)
260 | _xnorm += _norm
261 | _xnorm_cnt += 1
262 | _xnorm /= _xnorm_cnt
263 |
264 | embeddings = embeddings_list[0].copy()
265 | embeddings = sklearn.preprocessing.normalize(embeddings)
266 | acc1 = 0.0
267 | std1 = 0.0
268 | embeddings = embeddings_list[0] + embeddings_list[1]
269 | embeddings = sklearn.preprocessing.normalize(embeddings)
270 | print(embeddings.shape)
271 | print('infer time', time_consumed)
272 | _, _, accuracy, val, val_std, far = evaluate(embeddings, issame_list, nrof_folds=nfolds)
273 | acc2, std2 = np.mean(accuracy), np.std(accuracy)
274 | return acc1, std1, acc2, std2, _xnorm, embeddings_list
275 |
276 |
277 | def dumpR(data_set,
278 | backbone,
279 | batch_size,
280 | name='',
281 | data_extra=None,
282 | label_shape=None):
283 | print('dump verification embedding..')
284 | data_list = data_set[0]
285 | issame_list = data_set[1]
286 | embeddings_list = []
287 | time_consumed = 0.0
288 | for i in range(len(data_list)):
289 | data = data_list[i]
290 | embeddings = None
291 | ba = 0
292 | while ba < data.shape[0]:
293 | bb = min(ba + batch_size, data.shape[0])
294 | count = bb - ba
295 |
296 | _data = nd.slice_axis(data, axis=0, begin=bb - batch_size, end=bb)
297 | time0 = datetime.datetime.now()
298 | if data_extra is None:
299 | db = mx.io.DataBatch(data=(_data,), label=(_label,))
300 | else:
301 | db = mx.io.DataBatch(data=(_data, _data_extra),
302 | label=(_label,))
303 | model.forward(db, is_train=False)
304 | net_out = model.get_outputs()
305 | _embeddings = net_out[0].asnumpy()
306 | time_now = datetime.datetime.now()
307 | diff = time_now - time0
308 | time_consumed += diff.total_seconds()
309 | if embeddings is None:
310 | embeddings = np.zeros((data.shape[0], _embeddings.shape[1]))
311 | embeddings[ba:bb, :] = _embeddings[(batch_size - count):, :]
312 | ba = bb
313 | embeddings_list.append(embeddings)
314 | embeddings = embeddings_list[0] + embeddings_list[1]
315 | embeddings = sklearn.preprocessing.normalize(embeddings)
316 | actual_issame = np.asarray(issame_list)
317 | outname = os.path.join('temp.bin')
318 | with open(outname, 'wb') as f:
319 | pickle.dump((embeddings, issame_list),
320 | f,
321 | protocol=pickle.HIGHEST_PROTOCOL)
322 |
323 |
324 | # if __name__ == '__main__':
325 | #
326 | # parser = argparse.ArgumentParser(description='do verification')
327 | # # general
328 | # parser.add_argument('--data-dir', default='', help='')
329 | # parser.add_argument('--model',
330 | # default='../model/softmax,50',
331 | # help='path to load model.')
332 | # parser.add_argument('--target',
333 | # default='lfw,cfp_ff,cfp_fp,agedb_30',
334 | # help='test targets.')
335 | # parser.add_argument('--gpu', default=0, type=int, help='gpu id')
336 | # parser.add_argument('--batch-size', default=32, type=int, help='')
337 | # parser.add_argument('--max', default='', type=str, help='')
338 | # parser.add_argument('--mode', default=0, type=int, help='')
339 | # parser.add_argument('--nfolds', default=10, type=int, help='')
340 | # args = parser.parse_args()
341 | # image_size = [112, 112]
342 | # print('image_size', image_size)
343 | # ctx = mx.gpu(args.gpu)
344 | # nets = []
345 | # vec = args.model.split(',')
346 | # prefix = args.model.split(',')[0]
347 | # epochs = []
348 | # if len(vec) == 1:
349 | # pdir = os.path.dirname(prefix)
350 | # for fname in os.listdir(pdir):
351 | # if not fname.endswith('.params'):
352 | # continue
353 | # _file = os.path.join(pdir, fname)
354 | # if _file.startswith(prefix):
355 | # epoch = int(fname.split('.')[0].split('-')[1])
356 | # epochs.append(epoch)
357 | # epochs = sorted(epochs, reverse=True)
358 | # if len(args.max) > 0:
359 | # _max = [int(x) for x in args.max.split(',')]
360 | # assert len(_max) == 2
361 | # if len(epochs) > _max[1]:
362 | # epochs = epochs[_max[0]:_max[1]]
363 | #
364 | # else:
365 | # epochs = [int(x) for x in vec[1].split('|')]
366 | # print('model number', len(epochs))
367 | # time0 = datetime.datetime.now()
368 | # for epoch in epochs:
369 | # print('loading', prefix, epoch)
370 | # sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
371 | # # arg_params, aux_params = ch_dev(arg_params, aux_params, ctx)
372 | # all_layers = sym.get_internals()
373 | # sym = all_layers['fc1_output']
374 | # model = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
375 | # # model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
376 | # model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0],
377 | # image_size[1]))])
378 | # model.set_params(arg_params, aux_params)
379 | # nets.append(model)
380 | # time_now = datetime.datetime.now()
381 | # diff = time_now - time0
382 | # print('model loading time', diff.total_seconds())
383 | #
384 | # ver_list = []
385 | # ver_name_list = []
386 | # for name in args.target.split(','):
387 | # path = os.path.join(args.data_dir, name + ".bin")
388 | # if os.path.exists(path):
389 | # print('loading.. ', name)
390 | # data_set = load_bin(path, image_size)
391 | # ver_list.append(data_set)
392 | # ver_name_list.append(name)
393 | #
394 | # if args.mode == 0:
395 | # for i in range(len(ver_list)):
396 | # results = []
397 | # for model in nets:
398 | # acc1, std1, acc2, std2, xnorm, embeddings_list = test(
399 | # ver_list[i], model, args.batch_size, args.nfolds)
400 | # print('[%s]XNorm: %f' % (ver_name_list[i], xnorm))
401 | # print('[%s]Accuracy: %1.5f+-%1.5f' % (ver_name_list[i], acc1, std1))
402 | # print('[%s]Accuracy-Flip: %1.5f+-%1.5f' % (ver_name_list[i], acc2, std2))
403 | # results.append(acc2)
404 | # print('Max of [%s] is %1.5f' % (ver_name_list[i], np.max(results)))
405 | # elif args.mode == 1:
406 | # raise ValueError
407 | # else:
408 | # model = nets[0]
409 | # dumpR(ver_list[0], model, args.batch_size, args.target)
--------------------------------------------------------------------------------
/evaluation/FaceModel.py:
--------------------------------------------------------------------------------
1 | import cv2
2 | import numpy as np
3 | from sklearn.preprocessing import normalize
4 |
5 |
6 | class FaceModel():
7 | def __init__(self,model_prefix, model_epoch, ctx_id=7 , backbone="iresnet50"):
8 | self.gpu_id=ctx_id
9 | self.image_size = (112, 112)
10 | self.model_prefix=model_prefix
11 | self.model_epoch=model_epoch
12 | self.model=self._get_model(ctx=ctx_id,image_size=self.image_size,prefix=self.model_prefix,epoch=self.model_epoch,layer='fc1', backbone=backbone)
13 | def _get_model(self, ctx, image_size, prefix, epoch, layer):
14 | pass
15 |
16 | def _getFeatureBlob(self,input_blob):
17 | pass
18 |
19 | def get_feature(self, image_path, color="BGR"):
20 | image = cv2.imread(image_path)
21 | image = cv2.resize(image, (112, 112))
22 | if (color == "RGB"):
23 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
24 | a = np.transpose(image, (2, 0, 1))
25 | input_blob = np.expand_dims(a, axis=0)
26 | emb=self._getFeatureBlob(input_blob)
27 | emb = normalize(emb.reshape(1, -1))
28 | return emb
29 |
30 | def get_batch_feature(self, image_path_list, batch_size=16, color="BGR"):
31 | count = 0
32 | num_batch = int(len(image_path_list) / batch_size)
33 | features = []
34 | quality_score=[]
35 | for i in range(0, len(image_path_list), batch_size):
36 |
37 | if count < num_batch:
38 | tmp_list = image_path_list[i : i+batch_size]
39 | else:
40 | tmp_list = image_path_list[i :]
41 | count += 1
42 |
43 | images = []
44 | for image_path in tmp_list:
45 | image = cv2.imread(image_path)
46 | image = cv2.resize(image, (112, 112))
47 | if (color=="RGB"):
48 | image=cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
49 | a = np.transpose(image, (2, 0, 1))
50 | images.append(a)
51 | input_blob = np.array(images)
52 |
53 | emb, qs = self._getFeatureBlob(input_blob)
54 | quality_score.append(qs)
55 | features.append(emb)
56 | #print("batch"+str(i))
57 | features = np.vstack(features)
58 | quality_score=np.vstack(quality_score)
59 | features = normalize(features)
60 | return features, quality_score
61 |
--------------------------------------------------------------------------------
/evaluation/QualityModel.py:
--------------------------------------------------------------------------------
1 | import os
2 |
3 | from backbones.iresnet import iresnet100, iresnet50
4 | from evaluation.FaceModel import FaceModel
5 | import torch
6 | class QualityModel(FaceModel):
7 | def __init__(self, model_prefix, model_epoch, gpu_id):
8 | super(QualityModel, self).__init__(model_prefix, model_epoch, gpu_id)
9 |
10 | def _get_model(self, ctx, image_size, prefix, epoch, layer, backbone):
11 | weight = torch.load(os.path.join(prefix,epoch+"backbone.pth"))
12 | if (backbone=="iresnet50"):
13 | backbone = iresnet50(num_features=512, qs=1, use_se=False).to(f"cuda:{ctx}")
14 | else:
15 | backbone = iresnet100(num_features=512, qs=1, use_se=False).to(f"cuda:{ctx}")
16 |
17 | backbone.load_state_dict(weight)
18 | model = torch.nn.DataParallel(backbone, device_ids=[ctx])
19 | model.eval()
20 | return model
21 |
22 | @torch.no_grad()
23 | def _getFeatureBlob(self,input_blob):
24 | imgs = torch.Tensor(input_blob).cuda()
25 | imgs.div_(255).sub_(0.5).div_(0.5)
26 | feat, qs = self.model(imgs)
27 | return feat.cpu().numpy(), qs.cpu().numpy()
28 |
--------------------------------------------------------------------------------
/evaluation/crerate_pair_xqlfw.py:
--------------------------------------------------------------------------------
1 | import os
2 | import cv2
3 |
4 | path = "data/XQLFW"
5 | outpath = "data/quality_data/XQLFW"
6 |
7 | dataset_name = path.split("/")[-1]
8 | rel_img_path = os.path.join(outpath.split("/")[-1], dataset_name, "images")
9 | outpath = os.path.join(outpath, dataset_name)
10 | if not os.path.exists(outpath):
11 | os.makedirs(outpath)
12 | os.makedirs(os.path.join(outpath, "images"))
13 |
14 | align_path = os.path.join(path, "xqlfw_aligned_112")
15 |
16 |
17 | def copy_img(person, img_id):
18 | img_name = f"{person}_{str(img_id).zfill(4)}.jpg"
19 | tmp_path = os.path.join(align_path, person, img_name)
20 | img = cv2.imread(tmp_path)
21 | img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
22 | cv2.imwrite(os.path.join(outpath, "images", img_name), img)
23 | return img_name
24 |
25 |
26 | def create_xqlfw_pairs(pairs_filename):
27 | """ reads xqlfw pairs.txt and creates pair_list.txt
28 | and image_path_list.txt of required format
29 | :param pairs_filename: path to pairs.txt
30 | """
31 | txt_file = open(os.path.join(outpath, "image_path_list.txt"), "w")
32 | pair_list = open(os.path.join(outpath, "pair_list.txt"), "w")
33 |
34 | f = open(pairs_filename, 'r')
35 | for line in f.readlines()[1:]:
36 | pair = line.strip().split()
37 | if len(pair) == 3:
38 | img_name1 = copy_img(pair[0], pair[1])
39 | img_name2 = copy_img(pair[0], pair[2])
40 | else:
41 | img_name1 = copy_img(pair[0], pair[1])
42 | img_name2 = copy_img(pair[2], pair[3])
43 |
44 | txt_file.write(os.path.join(rel_img_path, img_name1)+"\n")
45 | txt_file.write(os.path.join(rel_img_path, img_name2)+"\n")
46 | pair_list.write(f"{img_name1} {img_name2} {int(len(pair)==3)}\n")
47 |
48 |
49 | f.close()
50 | txt_file.close()
51 | pair_list.close()
52 |
53 | create_xqlfw_pairs("data/XQLFW/xqlfw_pairs.txt")
54 | print("XQLFW successfully extracted")
--------------------------------------------------------------------------------
/evaluation/extract_adience.py:
--------------------------------------------------------------------------------
1 | import os
2 | import cv2
3 | import numpy as np
4 | from tqdm import tqdm
5 | from align_trans import norm_crop
6 |
7 | from mtcnn import MTCNN
8 | import tensorflow as tf
9 |
10 | config = tf.compat.v1.ConfigProto(gpu_options =
11 | tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=0.8)#, device_count = {'GPU': 0}
12 | )
13 | config.gpu_options.allow_growth = True
14 | config.gpu_options.visible_device_list= '0'
15 | session = tf.compat.v1.Session(config=config)
16 | tf.compat.v1.keras.backend.set_session(session)
17 |
18 |
19 | imgs_path = "/data/maklemt/adience"
20 | outpath = "/data/maklemt/quality_data"
21 |
22 |
23 | def find_central_face(img, keypoints):
24 | # if multiple faces are detected, select the face most in the middle of the image
25 | mid_face_idx = 0
26 | if len(keypoints) > 1:
27 | img_mid_point = np.array([img.shape[1]//2, img.shape[0]//2]) # [x (width), y (height)]
28 | noses = np.array([keypoint['keypoints']['nose'] for keypoint in keypoints])
29 | distances = np.linalg.norm(noses - img_mid_point, axis=1) # calculate distance between nose and img mid point
30 | mid_face_idx = np.argmin(distances)
31 |
32 | facial5points = [keypoints[mid_face_idx]['keypoints']['left_eye'], keypoints[mid_face_idx]['keypoints']['right_eye'],
33 | keypoints[mid_face_idx]['keypoints']['nose'], keypoints[mid_face_idx]['keypoints']['mouth_left'],
34 | keypoints[mid_face_idx]['keypoints']['mouth_right']]
35 | return np.array(facial5points)
36 |
37 |
38 |
39 | detector = MTCNN(min_face_size=20, steps_threshold=[0.6, 0.7, 0.9], scale_factor=0.85)
40 | skipped_imgs = []
41 |
42 | dataset_name = imgs_path.split("/")[-1]
43 | rel_img_path = os.path.join(outpath.split("/")[-1], dataset_name, "images")
44 | outpath = os.path.join(outpath, dataset_name)
45 | if not os.path.exists(outpath):
46 | os.makedirs(outpath)
47 | os.makedirs(os.path.join(outpath, "images"))
48 |
49 | print("extract:", dataset_name)
50 |
51 | txt_file = open(os.path.join(outpath, "image_path_list.txt"), "w")
52 | img_files = os.listdir(imgs_path)
53 |
54 | for img_index, img_file in tqdm(enumerate(img_files), total=len(img_files)):
55 | if img_file.split('.')[-1] != "jpg":
56 | continue
57 | img_path = os.path.join(imgs_path, img_file)
58 | img = cv2.imread(img_path)
59 |
60 | keypoints = detector.detect_faces(img)
61 | if len(keypoints) == 0:
62 | skipped_imgs.append(img_file)
63 | continue
64 | facial5points = find_central_face(img, keypoints)
65 | warped_face = norm_crop(img, landmark=facial5points, createEvalDB=True)
66 |
67 | img = cv2.cvtColor(warped_face, cv2.COLOR_RGB2BGR)
68 | person = img_file.split('.')[1]
69 | new_img_name = f"p_{person}_img_{img_index}.jpg"
70 |
71 | cv2.imwrite(os.path.join(outpath, "images", new_img_name), img)
72 | txt_file.write(os.path.join(rel_img_path, new_img_name)+"\n")
73 |
74 | txt_file.close()
75 |
76 |
77 | print("creating pair list...")
78 | pair_list = open(os.path.join(outpath, "pair_list.txt"), "w")
79 | aligned_img_path = os.listdir(os.path.join(outpath, "images"))
80 |
81 | for img_index1, img_file1 in enumerate(aligned_img_path):
82 | if img_file1.split('.')[-1] != "jpg":
83 | continue
84 | person1 = img_file1.split('_')[1]
85 |
86 | for img_index2, img_file2 in enumerate(aligned_img_path[img_index1+1:]):
87 | if img_file2.split('.')[-1] != "jpg":
88 | continue
89 | person2 = img_file2.split('_')[1]
90 | genuine = person1 == person2
91 | pair_list.write(f"{img_file1.split('.')[0]} {img_file2.split('.')[0]} {int(genuine)}\n")
92 |
93 | pair_list.close()
94 | print("pair_list saved")
95 | print("No faces detected in:")
96 | print(skipped_imgs)
97 | print("Total amount of images with no detected face: ", len(skipped_imgs))
98 |
--------------------------------------------------------------------------------
/evaluation/getQualityScore.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import os
3 | import sys
4 |
5 | from evaluation.QualityModel import QualityModel
6 |
7 |
8 | def parse_arguments(argv):
9 | parser = argparse.ArgumentParser()
10 |
11 | parser.add_argument('--data-dir', type=str, default='./data',
12 | help='Root dir for evaluation dataset')
13 | parser.add_argument('--pairs', type=str, default='pairs.txt',
14 | help='lfw pairs.')
15 | parser.add_argument('--datasets', type=str, default='XQLFW',
16 | help='list of evaluation datasets (,) e.g. XQLFW, lfw,calfw,agedb_30,cfp_fp,cplfw,IJBC.')
17 | parser.add_argument('--gpu_id', type=int, default=0,
18 | help='GPU id.')
19 | parser.add_argument('--model_path', type=str, default="/home/fboutros/LearnableMargin/output/ResNet50-COSQSArcFace_SmothL1",
20 | help='path to pretrained evaluation.')
21 | parser.add_argument('--model_id', type=str, default="32572",
22 | help='digit number in backbone file name')
23 | parser.add_argument('--backbone', type=str, default="iresnet50",
24 | help=' iresnet100 or iresnet50 ')
25 | parser.add_argument('--score_file_name', type=str, default="quality_r50.txt",
26 | help='score file name, the file will be store in the same data dir')
27 | parser.add_argument('--color_channel', type=str, default="BGR",
28 | help='input image color channel, two option RGB or BGR')
29 | #
30 |
31 | return parser.parse_args(argv)
32 |
33 | def read_image_list(image_list_file, image_dir=''):
34 | image_lists = []
35 | with open(image_list_file) as f:
36 | absolute_list=f.readlines()
37 | for l in absolute_list:
38 | image_lists.append(os.path.join(image_dir, l.rstrip()))
39 | return image_lists, absolute_list
40 | def main(param):
41 | datasets=param.datasets.split(',')
42 | face_model=QualityModel(param.model_path,param.model_id, param.gpu_id, param.backbone)
43 | for dataset in datasets:
44 | root=os.path.join(param.data_dir)
45 | image_list, absolute_list=read_image_list(os.path.join(param.data_dir,'quality_data',dataset,'image_path_list.txt'), root)
46 | embedding, quality=face_model.get_batch_feature(image_list,batch_size=16, color=param.color_channel)
47 | if (os.path.isdir(os.path.join(param.data_dir,'quality_data',dataset))):
48 | os.makedirs(os.path.join(param.data_dir,'quality_data',dataset))
49 | quality_score=open(os.path.join(param.data_dir,'quality_data',dataset,param.score_file_name),"a")
50 | for i in range(len(quality)):
51 | quality_score.write(absolute_list[i].rstrip()+ " "+str(quality[i][0])+ "\n")
52 |
53 | if __name__ == '__main__':
54 | main(parse_arguments(sys.argv[1:]))
--------------------------------------------------------------------------------
/feature_extraction/backbones/activation.py:
--------------------------------------------------------------------------------
1 | import torch.nn as nn
2 | import torch.nn.functional as F
3 |
4 | import torch
5 |
6 | from inspect import isfunction
7 |
8 | class Identity(nn.Module):
9 | """
10 | Identity block.
11 | """
12 | def __init__(self):
13 | super(Identity, self).__init__()
14 |
15 | def forward(self, x):
16 | return x
17 |
18 | def __repr__(self):
19 | return '{name}()'.format(name=self.__class__.__name__)
20 | class HSigmoid(nn.Module):
21 | """
22 | Approximated sigmoid function, so-called hard-version of sigmoid from 'Searching for MobileNetV3,'
23 | https://arxiv.org/abs/1905.02244.
24 | """
25 | def forward(self, x):
26 | return F.relu6(x + 3.0, inplace=True) / 6.0
27 |
28 |
29 | class Swish(nn.Module):
30 | """
31 | Swish activation function from 'Searching for Activation Functions,' https://arxiv.org/abs/1710.05941.
32 | """
33 | def forward(self, x):
34 | return x * torch.sigmoid(x)
35 | class HSwish(nn.Module):
36 | """
37 | H-Swish activation function from 'Searching for MobileNetV3,' https://arxiv.org/abs/1905.02244.
38 |
39 | Parameters:
40 | ----------
41 | inplace : bool
42 | Whether to use inplace version of the module.
43 | """
44 | def __init__(self, inplace=False):
45 | super(HSwish, self).__init__()
46 | self.inplace = inplace
47 |
48 | def forward(self, x):
49 | return x * F.relu6(x + 3.0, inplace=self.inplace) / 6.0
50 |
51 |
52 | def get_activation_layer(activation,param):
53 | """
54 | Create activation layer from string/function.
55 |
56 | Parameters:
57 | ----------
58 | activation : function, or str, or nn.Module
59 | Activation function or name of activation function.
60 |
61 | Returns:
62 | -------
63 | nn.Module
64 | Activation layer.
65 | """
66 | assert (activation is not None)
67 | if isfunction(activation):
68 | return activation()
69 | elif isinstance(activation, str):
70 | if activation == "relu":
71 | return nn.ReLU(inplace=True)
72 | elif activation =="prelu":
73 | return nn.PReLU(param)
74 | elif activation == "relu6":
75 | return nn.ReLU6(inplace=True)
76 | elif activation == "swish":
77 | return Swish()
78 | elif activation == "hswish":
79 | return HSwish(inplace=True)
80 | elif activation == "sigmoid":
81 | return nn.Sigmoid()
82 | elif activation == "hsigmoid":
83 | return HSigmoid()
84 | elif activation == "identity":
85 | return Identity()
86 | else:
87 | raise NotImplementedError()
88 | else:
89 | assert (isinstance(activation, nn.Module))
90 | return activation
--------------------------------------------------------------------------------
/feature_extraction/backbones/iresnet.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch import nn
3 |
4 | __all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100']
5 |
6 |
7 | def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
8 | """3x3 convolution with padding"""
9 | return nn.Conv2d(in_planes,
10 | out_planes,
11 | kernel_size=3,
12 | stride=stride,
13 | padding=dilation,
14 | groups=groups,
15 | bias=False,
16 | dilation=dilation)
17 |
18 |
19 | def conv1x1(in_planes, out_planes, stride=1):
20 | """1x1 convolution"""
21 | return nn.Conv2d(in_planes,
22 | out_planes,
23 | kernel_size=1,
24 | stride=stride,
25 | bias=False)
26 |
27 |
28 | class IBasicBlock(nn.Module):
29 | expansion = 1
30 | def __init__(self, inplanes, planes, stride=1, downsample=None,
31 | groups=1, base_width=64, dilation=1):
32 | super(IBasicBlock, self).__init__()
33 | if groups != 1 or base_width != 64:
34 | raise ValueError('BasicBlock only supports groups=1 and base_width=64')
35 | if dilation > 1:
36 | raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
37 | self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,)
38 | self.conv1 = conv3x3(inplanes, planes)
39 | self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,)
40 | self.prelu = nn.PReLU(planes)
41 | self.conv2 = conv3x3(planes, planes, stride)
42 | self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,)
43 | self.downsample = downsample
44 | self.stride = stride
45 |
46 | def forward(self, x):
47 | identity = x
48 | out = self.bn1(x)
49 | out = self.conv1(out)
50 | out = self.bn2(out)
51 | out = self.prelu(out)
52 | out = self.conv2(out)
53 | out = self.bn3(out)
54 | if self.downsample is not None:
55 | identity = self.downsample(x)
56 | out += identity
57 | return out
58 |
59 |
60 | class IResNet(nn.Module):
61 | fc_scale = 7 * 7
62 | def __init__(self,
63 | block, layers, dropout=0, num_features=512, zero_init_residual=False,
64 | groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False):
65 | super(IResNet, self).__init__()
66 | self.fp16 = fp16
67 | self.inplanes = 64
68 | self.dilation = 1
69 | if replace_stride_with_dilation is None:
70 | replace_stride_with_dilation = [False, False, False]
71 | if len(replace_stride_with_dilation) != 3:
72 | raise ValueError("replace_stride_with_dilation should be None "
73 | "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
74 | self.groups = groups
75 | self.base_width = width_per_group
76 | self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
77 | self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
78 | self.prelu = nn.PReLU(self.inplanes)
79 | self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
80 | self.layer2 = self._make_layer(block,
81 | 128,
82 | layers[1],
83 | stride=2,
84 | dilate=replace_stride_with_dilation[0])
85 | self.layer3 = self._make_layer(block,
86 | 256,
87 | layers[2],
88 | stride=2,
89 | dilate=replace_stride_with_dilation[1])
90 | self.layer4 = self._make_layer(block,
91 | 512,
92 | layers[3],
93 | stride=2,
94 | dilate=replace_stride_with_dilation[2])
95 | self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,)
96 | self.dropout = nn.Dropout(p=dropout, inplace=True)
97 | self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)
98 | self.features = nn.BatchNorm1d(num_features, eps=1e-05)
99 | nn.init.constant_(self.features.weight, 1.0)
100 | self.features.weight.requires_grad = False
101 |
102 | for m in self.modules():
103 | if isinstance(m, nn.Conv2d):
104 | nn.init.normal_(m.weight, 0, 0.1)
105 | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
106 | nn.init.constant_(m.weight, 1)
107 | nn.init.constant_(m.bias, 0)
108 |
109 | if zero_init_residual:
110 | for m in self.modules():
111 | if isinstance(m, IBasicBlock):
112 | nn.init.constant_(m.bn2.weight, 0)
113 |
114 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
115 | downsample = None
116 | previous_dilation = self.dilation
117 | if dilate:
118 | self.dilation *= stride
119 | stride = 1
120 | if stride != 1 or self.inplanes != planes * block.expansion:
121 | downsample = nn.Sequential(
122 | conv1x1(self.inplanes, planes * block.expansion, stride),
123 | nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),
124 | )
125 | layers = []
126 | layers.append(
127 | block(self.inplanes, planes, stride, downsample, self.groups,
128 | self.base_width, previous_dilation))
129 | self.inplanes = planes * block.expansion
130 | for _ in range(1, blocks):
131 | layers.append(
132 | block(self.inplanes,
133 | planes,
134 | groups=self.groups,
135 | base_width=self.base_width,
136 | dilation=self.dilation))
137 |
138 | return nn.Sequential(*layers)
139 |
140 | def forward(self, x):
141 | with torch.cuda.amp.autocast(self.fp16):
142 | x = self.conv1(x)
143 | x = self.bn1(x)
144 | x = self.prelu(x)
145 | x = self.layer1(x)
146 | x = self.layer2(x)
147 | x = self.layer3(x)
148 | x = self.layer4(x)
149 | x = self.bn2(x)
150 | x = torch.flatten(x, 1)
151 | x = self.dropout(x)
152 | x = self.fc(x.float() if self.fp16 else x)
153 | x = self.features(x)
154 | return x
155 |
156 |
157 | def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
158 | model = IResNet(block, layers, **kwargs)
159 | if pretrained:
160 | raise ValueError()
161 | return model
162 |
163 |
164 | def iresnet18(pretrained=False, progress=True, **kwargs):
165 | return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained,
166 | progress, **kwargs)
167 |
168 |
169 | def iresnet34(pretrained=False, progress=True, **kwargs):
170 | return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained,
171 | progress, **kwargs)
172 |
173 |
174 | def iresnet50(pretrained=False, progress=True, **kwargs):
175 | return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,
176 | progress, **kwargs)
177 |
178 |
179 | def iresnet100(pretrained=False, progress=True, **kwargs):
180 | return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained,
181 | progress, **kwargs)
182 | def _test():
183 | import torch
184 |
185 | pretrained = False
186 |
187 | models = [
188 | iresnet100
189 | ]
190 |
191 | for model in models:
192 |
193 | net = model()
194 | print(net)
195 | # net.train()
196 | weight_count = _calc_width(net)
197 | flops=count_model_flops(net)
198 | print("m={}, {}".format(model.__name__, weight_count))
199 | print("m={}, {}".format(model.__name__, flops))
200 |
201 | #assert (model != mixnet_s or weight_count == 4134606)
202 | #assert (model != mixnet_m or weight_count == 5014382)
203 | #assert (model != mixnet_l or weight_count == 7329252)
204 | net.eval()
205 |
206 | x = torch.randn(1, 3, 112, 112)
207 |
208 | y = net(x)
209 | y.sum().backward()
210 | assert (tuple(y.size()) == (1, 512))
211 |
212 |
213 | if __name__ == "__main__":
214 | _test()
215 |
--------------------------------------------------------------------------------
/feature_extraction/backbones/iresnet_mag.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch import nn
3 | # from torchvision.models.utils import load_state_dict_from_url
4 |
5 | __all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100']
6 |
7 |
8 | def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
9 | """3x3 convolution with padding"""
10 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
11 | padding=dilation, groups=groups, bias=False, dilation=dilation)
12 |
13 |
14 | def conv1x1(in_planes, out_planes, stride=1):
15 | """1x1 convolution"""
16 | return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
17 |
18 |
19 | class IBasicBlock(nn.Module):
20 | expansion = 1
21 |
22 | def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
23 | base_width=64, dilation=1):
24 | super(IBasicBlock, self).__init__()
25 | if groups != 1 or base_width != 64:
26 | raise ValueError(
27 | 'BasicBlock only supports groups=1 and base_width=64')
28 | if dilation > 1:
29 | raise NotImplementedError(
30 | "Dilation > 1 not supported in BasicBlock")
31 | # Both self.conv1 and self.downsample layers downsample the input when stride != 1
32 | self.bn1 = nn.BatchNorm2d(inplanes, eps=2e-05, momentum=0.9)
33 | self.conv1 = conv3x3(inplanes, planes)
34 | self.bn2 = nn.BatchNorm2d(planes, eps=2e-05, momentum=0.9)
35 | self.prelu = nn.PReLU(planes)
36 | self.conv2 = conv3x3(planes, planes, stride)
37 | self.bn3 = nn.BatchNorm2d(planes, eps=2e-05, momentum=0.9)
38 | self.downsample = downsample
39 | self.stride = stride
40 |
41 | def forward(self, x):
42 | identity = x
43 |
44 | out = self.bn1(x)
45 | out = self.conv1(out)
46 | out = self.bn2(out)
47 | out = self.prelu(out)
48 | out = self.conv2(out)
49 | out = self.bn3(out)
50 |
51 | if self.downsample is not None:
52 | identity = self.downsample(x)
53 |
54 | out += identity
55 |
56 | return out
57 |
58 |
59 | class IResNet(nn.Module):
60 | fc_scale = 7 * 7
61 |
62 | def __init__(self, block, layers, num_classes=512, zero_init_residual=False,
63 | groups=1, width_per_group=64, replace_stride_with_dilation=None):
64 | super(IResNet, self).__init__()
65 |
66 | self.inplanes = 64
67 | self.dilation = 1
68 | if replace_stride_with_dilation is None:
69 | # each element in the tuple indicates if we should replace
70 | # the 2x2 stride with a dilated convolution instead
71 | replace_stride_with_dilation = [False, False, False]
72 | if len(replace_stride_with_dilation) != 3:
73 | raise ValueError("replace_stride_with_dilation should be None "
74 | "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
75 | self.groups = groups
76 | self.base_width = width_per_group
77 | self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1,
78 | bias=False)
79 | self.bn1 = nn.BatchNorm2d(self.inplanes, eps=2e-05, momentum=0.9)
80 | self.prelu = nn.PReLU(self.inplanes)
81 | self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
82 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
83 | dilate=replace_stride_with_dilation[0])
84 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
85 | dilate=replace_stride_with_dilation[1])
86 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
87 | dilate=replace_stride_with_dilation[2])
88 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
89 |
90 | self.bn2 = nn.BatchNorm2d(
91 | 512 * block.expansion, eps=2e-05, momentum=0.9)
92 | self.dropout = nn.Dropout2d(p=0.4, inplace=True)
93 | self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_classes)
94 | self.features = nn.BatchNorm1d(num_classes, eps=2e-05, momentum=0.9)
95 |
96 | for m in self.modules():
97 | if isinstance(m, nn.Conv2d):
98 | nn.init.kaiming_normal_(
99 | m.weight, mode='fan_out', nonlinearity='relu')
100 | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
101 | nn.init.constant_(m.weight, 1)
102 | nn.init.constant_(m.bias, 0)
103 |
104 | if zero_init_residual:
105 | for m in self.modules():
106 | if isinstance(m, IBasicBlock):
107 | nn.init.constant_(m.bn2.weight, 0)
108 |
109 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
110 | downsample = None
111 | previous_dilation = self.dilation
112 | if dilate:
113 | self.dilation *= stride
114 | stride = 1
115 | if stride != 1 or self.inplanes != planes * block.expansion:
116 | downsample = nn.Sequential(
117 | conv1x1(self.inplanes, planes * block.expansion, stride),
118 | nn.BatchNorm2d(planes * block.expansion,
119 | eps=2e-05, momentum=0.9),
120 | )
121 |
122 | layers = []
123 | layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
124 | self.base_width, previous_dilation))
125 | self.inplanes = planes * block.expansion
126 | for _ in range(1, blocks):
127 | layers.append(block(self.inplanes, planes, groups=self.groups,
128 | base_width=self.base_width, dilation=self.dilation))
129 |
130 | return nn.Sequential(*layers)
131 |
132 | def forward(self, x):
133 | x = self.conv1(x)
134 | x = self.bn1(x)
135 | x = self.prelu(x)
136 |
137 | x = self.layer1(x)
138 | x = self.layer2(x)
139 | x = self.layer3(x)
140 | x = self.layer4(x)
141 |
142 | x = self.bn2(x)
143 | x = self.dropout(x)
144 | x = x.view(x.size(0), -1)
145 | x = self.fc(x)
146 | x = self.features(x)
147 |
148 | return x
149 |
150 |
151 | def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
152 | model = IResNet(block, layers, **kwargs)
153 | # if pretrained:
154 | # state_dict = load_state_dict_from_url(model_urls[arch],
155 | # progress=progress)
156 | # model.load_state_dict(state_dict)
157 | return model
158 |
159 |
160 | def iresnet18(pretrained=False, progress=True, **kwargs):
161 | return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained, progress,
162 | **kwargs)
163 |
164 |
165 | def iresnet34(pretrained=False, progress=True, **kwargs):
166 | return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained, progress,
167 | **kwargs)
168 |
169 |
170 | def iresnet50(pretrained=False, progress=True, **kwargs):
171 | return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained, progress,
172 | **kwargs)
173 |
174 |
175 | def iresnet100(pretrained=False, progress=True, **kwargs):
176 | return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained, progress,
177 | **kwargs)
178 |
--------------------------------------------------------------------------------
/feature_extraction/backbones/mag_network_inf.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python
2 | import sys
3 | sys.path.append("..")
4 | from backbones import iresnet_mag as iresnet
5 | from collections import OrderedDict
6 | import os
7 | import torch.nn.functional as F
8 | import torch.nn as nn
9 | import torch
10 |
11 |
12 | def load_features(backbone_name, embedding_size):
13 | if backbone_name == 'iresnet34':
14 | features = iresnet.iresnet34(
15 | pretrained=False,
16 | num_classes=embedding_size,
17 | )
18 | elif backbone_name == 'iresnet18':
19 | features = iresnet.iresnet18(
20 | pretrained=False,
21 | num_classes=embedding_size,
22 | )
23 | elif backbone_name == 'iresnet50':
24 | features = iresnet.iresnet50(
25 | pretrained=False,
26 | num_classes=embedding_size,
27 | )
28 | elif backbone_name == 'iresnet100':
29 | features = iresnet.iresnet100(
30 | pretrained=False,
31 | num_classes=embedding_size,
32 | )
33 | else:
34 | raise ValueError()
35 | return features
36 |
37 |
38 | class NetworkBuilder_inf(nn.Module):
39 | def __init__(self, backbone, embedding_size):
40 | super(NetworkBuilder_inf, self).__init__()
41 | self.features = load_features(backbone, embedding_size)
42 |
43 | def forward(self, input):
44 | # add Fp, a pose feature
45 | x = self.features(input)
46 | return x
47 |
48 |
49 | def load_dict_inf(checkpoint_path, model, cpu_mode=False):
50 | if os.path.isfile(checkpoint_path):
51 | print('=> loading pth from {} ...'.format(checkpoint_path))
52 | if cpu_mode:
53 | checkpoint = torch.load(checkpoint_path, map_location=torch.device("cpu"))
54 | else:
55 | checkpoint = torch.load(checkpoint_path)
56 | _state_dict = clean_dict_inf(model, checkpoint['state_dict'])
57 | model_dict = model.state_dict()
58 | model_dict.update(_state_dict)
59 | model.load_state_dict(model_dict)
60 | # delete to release more space
61 | del checkpoint
62 | del _state_dict
63 | else:
64 | sys.exit("=> No checkpoint found at '{}'".format(checkpoint_path))
65 | return model
66 |
67 |
68 | def clean_dict_inf(model, state_dict):
69 | _state_dict = OrderedDict()
70 | for k, v in state_dict.items():
71 | # # assert k[0:1] == 'features.module.'
72 | new_k = 'features.'+'.'.join(k.split('.')[2:])
73 | if new_k in model.state_dict().keys() and \
74 | v.size() == model.state_dict()[new_k].size():
75 | _state_dict[new_k] = v
76 | # assert k[0:1] == 'module.features.'
77 | new_kk = '.'.join(k.split('.')[1:])
78 | if new_kk in model.state_dict().keys() and \
79 | v.size() == model.state_dict()[new_kk].size():
80 | _state_dict[new_kk] = v
81 | num_model = len(model.state_dict().keys())
82 | num_ckpt = len(_state_dict.keys())
83 | if num_model != num_ckpt:
84 | sys.exit("=> Not all weights loaded, model params: {}, loaded params: {}".format(
85 | num_model, num_ckpt))
86 | return _state_dict
87 |
88 |
89 | def builder_inf(checkpoint_path, backbone, embedding_size):
90 | model = NetworkBuilder_inf(backbone, embedding_size)
91 | # Used to run inference
92 | model = load_dict_inf(checkpoint_path, model)
93 | return model
94 |
--------------------------------------------------------------------------------
/feature_extraction/backbones/model_irse.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, ReLU, Sigmoid, Dropout, MaxPool2d, \
4 | AdaptiveAvgPool2d, Sequential, Module
5 | from collections import namedtuple
6 |
7 |
8 | # Support: ['IR_50', 'IR_101', 'IR_152', 'IR_SE_50', 'IR_SE_101', 'IR_SE_152']
9 |
10 |
11 | class Flatten(Module):
12 | def forward(self, input):
13 | return input.view(input.size(0), -1)
14 |
15 |
16 | def l2_norm(input, axis=1):
17 | norm = torch.norm(input, 2, axis, True)
18 | output = torch.div(input, norm)
19 |
20 | return output
21 |
22 |
23 | class SEModule(Module):
24 | def __init__(self, channels, reduction):
25 | super(SEModule, self).__init__()
26 | self.avg_pool = AdaptiveAvgPool2d(1)
27 | self.fc1 = Conv2d(
28 | channels, channels // reduction, kernel_size=1, padding=0, bias=False)
29 |
30 | nn.init.xavier_uniform_(self.fc1.weight.data)
31 |
32 | self.relu = ReLU(inplace=True)
33 | self.fc2 = Conv2d(
34 | channels // reduction, channels, kernel_size=1, padding=0, bias=False)
35 |
36 | self.sigmoid = Sigmoid()
37 |
38 | def forward(self, x):
39 | module_input = x
40 | x = self.avg_pool(x)
41 | x = self.fc1(x)
42 | x = self.relu(x)
43 | x = self.fc2(x)
44 | x = self.sigmoid(x)
45 |
46 | return module_input * x
47 |
48 |
49 | class bottleneck_IR(Module):
50 | def __init__(self, in_channel, depth, stride):
51 | super(bottleneck_IR, self).__init__()
52 | if in_channel == depth:
53 | self.shortcut_layer = MaxPool2d(1, stride)
54 | else:
55 | self.shortcut_layer = Sequential(
56 | Conv2d(in_channel, depth, (1, 1), stride, bias=False), BatchNorm2d(depth))
57 | self.res_layer = Sequential(
58 | BatchNorm2d(in_channel),
59 | Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
60 | PReLU(depth),
61 | Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
62 | BatchNorm2d(depth))
63 |
64 | def forward(self, x):
65 | shortcut = self.shortcut_layer(x)
66 | res = self.res_layer(x)
67 |
68 | return res + shortcut
69 |
70 |
71 | class bottleneck_IR_SE(Module):
72 | def __init__(self, in_channel, depth, stride):
73 | super(bottleneck_IR_SE, self).__init__()
74 | if in_channel == depth:
75 | self.shortcut_layer = MaxPool2d(1, stride)
76 | else:
77 | self.shortcut_layer = Sequential(
78 | Conv2d(in_channel, depth, (1, 1), stride, bias=False),
79 | BatchNorm2d(depth))
80 | self.res_layer = Sequential(
81 | BatchNorm2d(in_channel),
82 | Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
83 | PReLU(depth),
84 | Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
85 | BatchNorm2d(depth),
86 | SEModule(depth, 16)
87 | )
88 |
89 | def forward(self, x):
90 | shortcut = self.shortcut_layer(x)
91 | res = self.res_layer(x)
92 |
93 | return res + shortcut
94 |
95 |
96 | class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
97 | '''A named tuple describing a ResNet block.'''
98 |
99 |
100 | def get_block(in_channel, depth, num_units, stride=2):
101 |
102 | return [Bottleneck(in_channel, depth, stride)] + [Bottleneck(depth, depth, 1) for i in range(num_units - 1)]
103 |
104 |
105 | def get_blocks(num_layers):
106 | if num_layers == 50:
107 | blocks = [
108 | get_block(in_channel=64, depth=64, num_units=3),
109 | get_block(in_channel=64, depth=128, num_units=4),
110 | get_block(in_channel=128, depth=256, num_units=14),
111 | get_block(in_channel=256, depth=512, num_units=3)
112 | ]
113 | elif num_layers == 100:
114 | blocks = [
115 | get_block(in_channel=64, depth=64, num_units=3),
116 | get_block(in_channel=64, depth=128, num_units=13),
117 | get_block(in_channel=128, depth=256, num_units=30),
118 | get_block(in_channel=256, depth=512, num_units=3)
119 | ]
120 | elif num_layers == 152:
121 | blocks = [
122 | get_block(in_channel=64, depth=64, num_units=3),
123 | get_block(in_channel=64, depth=128, num_units=8),
124 | get_block(in_channel=128, depth=256, num_units=36),
125 | get_block(in_channel=256, depth=512, num_units=3)
126 | ]
127 |
128 | return blocks
129 |
130 |
131 | class Backbone(Module):
132 | def __init__(self, input_size, num_layers, mode='ir'):
133 | super(Backbone, self).__init__()
134 | assert input_size[0] in [112, 224], "input_size should be [112, 112] or [224, 224]"
135 | assert num_layers in [50, 100, 152], "num_layers should be 50, 100 or 152"
136 | assert mode in ['ir', 'ir_se'], "mode should be ir or ir_se"
137 | blocks = get_blocks(num_layers)
138 | if mode == 'ir':
139 | unit_module = bottleneck_IR
140 | elif mode == 'ir_se':
141 | unit_module = bottleneck_IR_SE
142 | self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
143 | BatchNorm2d(64),
144 | PReLU(64))
145 | if input_size[0] == 112:
146 | self.output_layer = Sequential(BatchNorm2d(512),
147 | Dropout(0.4),
148 | Flatten(),
149 | Linear(512 * 7 * 7, 512),
150 | BatchNorm1d(512, affine=False))
151 | else:
152 | self.output_layer = Sequential(BatchNorm2d(512),
153 | Dropout(0.4),
154 | Flatten(),
155 | Linear(512 * 14 * 14, 512),
156 | BatchNorm1d(512, affine=False))
157 |
158 | modules = []
159 | for block in blocks:
160 | for bottleneck in block:
161 | modules.append(
162 | unit_module(bottleneck.in_channel,
163 | bottleneck.depth,
164 | bottleneck.stride))
165 | self.body = Sequential(*modules)
166 |
167 | self._initialize_weights()
168 |
169 | def forward(self, x):
170 | x = self.input_layer(x)
171 | x = self.body(x)
172 | conv_out = x.view(x.shape[0], -1)
173 | x = self.output_layer(x)
174 |
175 | return x, conv_out
176 |
177 | def _initialize_weights(self):
178 | for m in self.modules():
179 | if isinstance(m, nn.Conv2d):
180 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
181 | if m.bias is not None:
182 | m.bias.data.zero_()
183 | elif isinstance(m, nn.BatchNorm2d):
184 | m.weight.data.fill_(1)
185 | m.bias.data.zero_()
186 | elif isinstance(m, nn.Linear):
187 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
188 | if m.bias is not None:
189 | m.bias.data.zero_()
190 |
191 | def IR_50(input_size):
192 | """Constructs a ir-50 model.
193 | """
194 | model = Backbone(input_size, 50, 'ir')
195 |
196 | return model
197 |
198 |
199 | def IR_101(input_size):
200 | """Constructs a ir-101 model.
201 | """
202 | model = Backbone(input_size, 100, 'ir')
203 |
204 | return model
205 |
206 |
207 | def IR_152(input_size):
208 | """Constructs a ir-152 model.
209 | """
210 | model = Backbone(input_size, 152, 'ir')
211 |
212 | return model
213 |
214 |
215 | def IR_SE_50(input_size):
216 | """Constructs a ir_se-50 model.
217 | """
218 | model = Backbone(input_size, 50, 'ir_se')
219 |
220 | return model
221 |
222 |
223 | def IR_SE_101(input_size):
224 | """Constructs a ir_se-101 model.
225 | """
226 | model = Backbone(input_size, 100, 'ir_se')
227 |
228 | return model
229 |
230 |
231 | def IR_SE_152(input_size):
232 | """Constructs a ir_se-152 model.
233 | """
234 | model = Backbone(input_size, 152, 'ir_se')
235 |
236 | return model
--------------------------------------------------------------------------------
/feature_extraction/backbones/utils.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch import nn
3 |
4 | from backbones.activation import get_activation_layer, HSwish, Swish
5 | from torch.autograd import Variable
6 |
7 | import numpy as np
8 |
9 | def round_channels(channels,
10 | divisor=8):
11 | """
12 | Round weighted channel number (make divisible operation).
13 |
14 | Parameters:
15 | ----------
16 | channels : int or float
17 | Original number of channels.
18 | divisor : int, default 8
19 | Alignment value.
20 |
21 | Returns:
22 | -------
23 | int
24 | Weighted number of channels.
25 | """
26 | rounded_channels = max(int(channels + divisor / 2.0) // divisor * divisor, divisor)
27 | if float(rounded_channels) < 0.9 * channels:
28 | rounded_channels += divisor
29 | return rounded_channels
30 |
31 | def conv1x1(in_channels,
32 | out_channels,
33 | stride=1,
34 | groups=1, dilation=1,
35 | bias=False):
36 | """
37 | Convolution 1x1 layer.
38 |
39 | Parameters:
40 | ----------
41 | in_channels : int
42 | Number of input channels.
43 | out_channels : int
44 | Number of output channels.
45 | stride : int or tuple/list of 2 int, default 1
46 | Strides of the convolution.
47 | groups : int, default 1
48 | Number of groups.
49 | bias : bool, default False
50 | Whether the layer uses a bias vector.
51 | """
52 | return nn.Conv2d(
53 | in_channels=in_channels,
54 | out_channels=out_channels,
55 | kernel_size=1,
56 | stride=stride,
57 | groups=groups, dilation=dilation,
58 | bias=bias)
59 |
60 |
61 |
62 |
63 | def conv3x3(in_channels,
64 | out_channels,
65 | stride=1,
66 | padding=1,
67 | dilation=1,
68 | groups=1,
69 | bias=False):
70 | """
71 | Convolution 3x3 layer.
72 |
73 | Parameters:
74 | ----------
75 | in_channels : int
76 | Number of input channels.
77 | out_channels : int
78 | Number of output channels.
79 | stride : int or tuple/list of 2 int, default 1
80 | Strides of the convolution.
81 | padding : int or tuple/list of 2 int, default 1
82 | Padding value for convolution layer.
83 | dilation : int or tuple/list of 2 int, default 1
84 | Dilation value for convolution layer.
85 | groups : int, default 1
86 | Number of groups.
87 | bias : bool, default False
88 | Whether the layer uses a bias vector.
89 | """
90 | return nn.Conv2d(
91 | in_channels=in_channels,
92 | out_channels=out_channels,
93 | kernel_size=3,
94 | stride=stride,
95 | padding=padding,
96 | dilation=dilation,
97 | groups=groups,
98 | bias=bias)
99 | class Flatten(nn.Module):
100 | """
101 | Simple flatten module.
102 | """
103 |
104 | def forward(self, x):
105 | return x.view(x.size(0), -1)
106 |
107 | def depthwise_conv3x3(channels,
108 | stride=1,
109 | padding=1,
110 | dilation=1,
111 | bias=False):
112 | """
113 | Depthwise convolution 3x3 layer.
114 |
115 | Parameters:
116 | ----------
117 | channels : int
118 | Number of input/output channels.
119 | strides : int or tuple/list of 2 int, default 1
120 | Strides of the convolution.
121 | padding : int or tuple/list of 2 int, default 1
122 | Padding value for convolution layer.
123 | dilation : int or tuple/list of 2 int, default 1
124 | Dilation value for convolution layer.
125 | bias : bool, default False
126 | Whether the layer uses a bias vector.
127 | """
128 | return nn.Conv2d(
129 | in_channels=channels,
130 | out_channels=channels,
131 | kernel_size=3,
132 | stride=stride,
133 | padding=padding,
134 | dilation=dilation,
135 | groups=channels,
136 | bias=bias)
137 | class ConvBlock(nn.Module):
138 | """
139 | Standard convolution block with Batch normalization and activation.
140 |
141 | Parameters:
142 | ----------
143 | in_channels : int
144 | Number of input channels.
145 | out_channels : int
146 | Number of output channels.
147 | kernel_size : int or tuple/list of 2 int
148 | Convolution window size.
149 | stride : int or tuple/list of 2 int
150 | Strides of the convolution.
151 | padding : int, or tuple/list of 2 int, or tuple/list of 4 int
152 | Padding value for convolution layer.
153 | dilation : int or tuple/list of 2 int, default 1
154 | Dilation value for convolution layer.
155 | groups : int, default 1
156 | Number of groups.
157 | bias : bool, default False
158 | Whether the layer uses a bias vector.
159 | use_bn : bool, default True
160 | Whether to use BatchNorm layer.
161 | bn_eps : float, default 1e-5
162 | Small float added to variance in Batch norm.
163 | activation : function or str or None, default nn.ReLU(inplace=True)
164 | Activation function or name of activation function.
165 | """
166 | def __init__(self,
167 | in_channels,
168 | out_channels,
169 | kernel_size,
170 | stride,
171 | padding,
172 | dilation=1,
173 | groups=1,
174 | bias=False,
175 | use_bn=True,
176 | bn_eps=1e-5,
177 | activation=(lambda: nn.ReLU(inplace=True))):
178 | super(ConvBlock, self).__init__()
179 | self.activate = (activation is not None)
180 | self.use_bn = use_bn
181 | self.use_pad = (isinstance(padding, (list, tuple)) and (len(padding) == 4))
182 |
183 | if self.use_pad:
184 | self.pad = nn.ZeroPad2d(padding=padding)
185 | padding = 0
186 | self.conv = nn.Conv2d(
187 | in_channels=in_channels,
188 | out_channels=out_channels,
189 | kernel_size=kernel_size,
190 | stride=stride,
191 | padding=padding,
192 | dilation=dilation,
193 | groups=groups,
194 | bias=bias)
195 | if self.use_bn:
196 | self.bn = nn.BatchNorm2d(
197 | num_features=out_channels,
198 | eps=bn_eps)
199 | if self.activate:
200 | self.activ = get_activation_layer(activation,out_channels)
201 |
202 | def forward(self, x):
203 | if self.use_pad:
204 | x = self.pad(x)
205 | x = self.conv(x)
206 | if self.use_bn:
207 | x = self.bn(x)
208 | if self.activate:
209 | x = self.activ(x)
210 | return x
211 |
212 | def conv1x1_block(in_channels,
213 | out_channels,
214 | stride=1,
215 | padding=0,
216 | groups=1,
217 | bias=False,
218 | use_bn=True,
219 | bn_eps=1e-5,
220 | activation=(lambda: nn.ReLU(inplace=True))):
221 | """
222 | 1x1 version of the standard convolution block.
223 |
224 | Parameters:
225 | ----------
226 | in_channels : int
227 | Number of input channels.
228 | out_channels : int
229 | Number of output channels.
230 | stride : int or tuple/list of 2 int, default 1
231 | Strides of the convolution.
232 | padding : int, or tuple/list of 2 int, or tuple/list of 4 int, default 0
233 | Padding value for convolution layer.
234 | groups : int, default 1
235 | Number of groups.
236 | bias : bool, default False
237 | Whether the layer uses a bias vector.
238 | use_bn : bool, default True
239 | Whether to use BatchNorm layer.
240 | bn_eps : float, default 1e-5
241 | Small float added to variance in Batch norm.
242 | activation : function or str or None, default nn.ReLU(inplace=True)
243 | Activation function or name of activation function.
244 | """
245 | return ConvBlock(
246 | in_channels=in_channels,
247 | out_channels=out_channels,
248 | kernel_size=1,
249 | stride=stride,
250 | padding=padding,
251 | groups=groups,
252 | bias=bias,
253 | use_bn=use_bn,
254 | bn_eps=bn_eps,
255 | activation=activation)
256 |
257 | def conv3x3_block(in_channels,
258 | out_channels,
259 | stride=1,
260 | padding=1,
261 | dilation=1,
262 | groups=1,
263 | bias=False,
264 | use_bn=True,
265 | bn_eps=1e-5,
266 | activation=(lambda: nn.ReLU(inplace=True))):
267 | """
268 | 3x3 version of the standard convolution block.
269 |
270 | Parameters:
271 | ----------
272 | in_channels : int
273 | Number of input channels.
274 | out_channels : int
275 | Number of output channels.
276 | stride : int or tuple/list of 2 int, default 1
277 | Strides of the convolution.
278 | padding : int, or tuple/list of 2 int, or tuple/list of 4 int, default 1
279 | Padding value for convolution layer.
280 | dilation : int or tuple/list of 2 int, default 1
281 | Dilation value for convolution layer.
282 | groups : int, default 1
283 | Number of groups.
284 | bias : bool, default False
285 | Whether the layer uses a bias vector.
286 | use_bn : bool, default True
287 | Whether to use BatchNorm layer.
288 | bn_eps : float, default 1e-5
289 | Small float added to variance in Batch norm.
290 | activation : function or str or None, default nn.ReLU(inplace=True)
291 | Activation function or name of activation function.
292 | """
293 | return ConvBlock(
294 | in_channels=in_channels,
295 | out_channels=out_channels,
296 | kernel_size=3,
297 | stride=stride,
298 | padding=padding,
299 | dilation=dilation,
300 | groups=groups,
301 | bias=bias,
302 | use_bn=use_bn,
303 | bn_eps=bn_eps,
304 | activation=activation)
305 | class DwsConvBlock(nn.Module):
306 | """
307 | Depthwise separable convolution block with BatchNorms and activations at each convolution layers.
308 |
309 | Parameters:
310 | ----------
311 | in_channels : int
312 | Number of input channels.
313 | out_channels : int
314 | Number of output channels.
315 | kernel_size : int or tuple/list of 2 int
316 | Convolution window size.
317 | stride : int or tuple/list of 2 int
318 | Strides of the convolution.
319 | padding : int, or tuple/list of 2 int, or tuple/list of 4 int
320 | Padding value for convolution layer.
321 | dilation : int or tuple/list of 2 int, default 1
322 | Dilation value for convolution layer.
323 | bias : bool, default False
324 | Whether the layer uses a bias vector.
325 | dw_use_bn : bool, default True
326 | Whether to use BatchNorm layer (depthwise convolution block).
327 | pw_use_bn : bool, default True
328 | Whether to use BatchNorm layer (pointwise convolution block).
329 | bn_eps : float, default 1e-5
330 | Small float added to variance in Batch norm.
331 | dw_activation : function or str or None, default nn.ReLU(inplace=True)
332 | Activation function after the depthwise convolution block.
333 | pw_activation : function or str or None, default nn.ReLU(inplace=True)
334 | Activation function after the pointwise convolution block.
335 | """
336 | def __init__(self,
337 | in_channels,
338 | out_channels,
339 | kernel_size,
340 | stride,
341 | padding,
342 | dilation=1,
343 | bias=False,
344 | dw_use_bn=True,
345 | pw_use_bn=True,
346 | bn_eps=1e-5,
347 | dw_activation=(lambda: nn.ReLU(inplace=True)),
348 | pw_activation=(lambda: nn.ReLU(inplace=True))):
349 | super(DwsConvBlock, self).__init__()
350 | self.dw_conv = dwconv_block(
351 | in_channels=in_channels,
352 | out_channels=in_channels,
353 | kernel_size=kernel_size,
354 | stride=stride,
355 | padding=padding,
356 | dilation=dilation,
357 | bias=bias,
358 | use_bn=dw_use_bn,
359 | bn_eps=bn_eps,
360 | activation=dw_activation)
361 | self.pw_conv = conv1x1_block(
362 | in_channels=in_channels,
363 | out_channels=out_channels,
364 | bias=bias,
365 | use_bn=pw_use_bn,
366 | bn_eps=bn_eps,
367 | activation=pw_activation)
368 |
369 | def forward(self, x):
370 | x = self.dw_conv(x)
371 |
372 | x = self.pw_conv(x)
373 |
374 | return x
375 |
376 |
377 | def dwconv_block(in_channels,
378 | out_channels,
379 | kernel_size,
380 | stride=1,
381 | padding=1,
382 | dilation=1,
383 | bias=False,
384 | use_bn=True,
385 | bn_eps=1e-5,
386 | activation=(lambda: nn.ReLU(inplace=True))):
387 | """
388 | Depthwise convolution block.
389 | """
390 | return ConvBlock(
391 | in_channels=in_channels,
392 | out_channels=out_channels,
393 | kernel_size=kernel_size,
394 | stride=stride,
395 | padding=padding,
396 | dilation=dilation,
397 | groups=out_channels,
398 | bias=bias,
399 | use_bn=use_bn,
400 | bn_eps=bn_eps,
401 | activation=activation)
402 |
403 | def channel_shuffle2(x,
404 | groups):
405 | """
406 | Channel shuffle operation from 'ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices,'
407 | https://arxiv.org/abs/1707.01083. The alternative version.
408 |
409 | Parameters:
410 | ----------
411 | x : Tensor
412 | Input tensor.
413 | groups : int
414 | Number of groups.
415 |
416 | Returns:
417 | -------
418 | Tensor
419 | Resulted tensor.
420 | """
421 | batch, channels, height, width = x.size()
422 | assert (channels % groups == 0)
423 | channels_per_group = channels // groups
424 |
425 | x = x.view(batch, channels_per_group, groups, height, width)
426 | x = torch.transpose(x, 1, 2).contiguous()
427 |
428 | x = x.view(batch, channels, height, width)
429 | return x
430 |
431 | def _calc_width(net):
432 | import numpy as np
433 | net_params = filter(lambda p: p.requires_grad, net.parameters())
434 | weight_count = 0
435 | for param in net_params:
436 | weight_count += np.prod(param.size())
437 | return weight_count
438 |
439 | def flops_to_string(flops, units='GFLOPS', precision=4):
440 | if units == 'GFLOPS':
441 | return str(round(flops / 10.**9, precision)) + ' ' + units
442 | elif units == 'MFLOPS':
443 | return str(round(flops / 10.**6, precision)) + ' ' + units
444 | elif units == 'KFLOPS':
445 | return str(round(flops / 10.**3, precision)) + ' ' + units
446 | else:
447 | return str(flops) + ' FLOPS'
448 |
449 |
450 |
451 |
452 | def count_model_flops(model, input_res=[112, 112], multiply_adds=True):
453 | list_conv = []
454 |
455 | def conv_hook(self, input, output):
456 | batch_size, input_channels, input_height, input_width = input[0].size()
457 | output_channels, output_height, output_width = output[0].size()
458 |
459 | kernel_ops = self.kernel_size[0] * self.kernel_size[1] * (self.in_channels / self.groups)
460 | bias_ops = 1 if self.bias is not None else 0
461 |
462 | params = output_channels * (kernel_ops + bias_ops)
463 | flops = (kernel_ops * (
464 | 2 if multiply_adds else 1) + bias_ops) * output_channels * output_height * output_width * batch_size
465 |
466 | list_conv.append(flops)
467 |
468 | list_linear = []
469 |
470 | def linear_hook(self, input, output):
471 | batch_size = input[0].size(0) if input[0].dim() == 2 else 1
472 |
473 | weight_ops = self.weight.nelement() * (2 if multiply_adds else 1)
474 | if self.bias is not None:
475 | bias_ops = self.bias.nelement() if self.bias.nelement() else 0
476 | flops = batch_size * (weight_ops + bias_ops)
477 | else:
478 | flops = batch_size * weight_ops
479 | list_linear.append(flops)
480 |
481 | list_bn = []
482 |
483 | def bn_hook(self, input, output):
484 | list_bn.append(input[0].nelement() * 2)
485 |
486 | list_relu = []
487 |
488 | def relu_hook(self, input, output):
489 | list_relu.append(input[0].nelement())
490 |
491 | list_pooling = []
492 |
493 | def pooling_hook(self, input, output):
494 | batch_size, input_channels, input_height, input_width = input[0].size()
495 | output_channels, output_height, output_width = output[0].size()
496 |
497 | kernel_ops = self.kernel_size * self.kernel_size
498 | bias_ops = 0
499 | params = 0
500 | flops = (kernel_ops + bias_ops) * output_channels * output_height * output_width * batch_size
501 |
502 | list_pooling.append(flops)
503 | def pooling_hook_ad(self, input, output):
504 | batch_size, input_channels, input_height, input_width = input[0].size()
505 | input = input[0]
506 | flops = int(np.prod(input.shape))
507 | list_pooling.append(flops)
508 |
509 | handles = []
510 |
511 | def foo(net):
512 | childrens = list(net.children())
513 | if not childrens:
514 | if isinstance(net, torch.nn.Conv2d) or isinstance(net, torch.nn.ConvTranspose2d):
515 | handles.append(net.register_forward_hook(conv_hook))
516 | elif isinstance(net, torch.nn.Linear):
517 | handles.append(net.register_forward_hook(linear_hook))
518 | elif isinstance(net, torch.nn.BatchNorm2d) or isinstance(net, torch.nn.BatchNorm1d):
519 | handles.append(net.register_forward_hook(bn_hook))
520 | elif isinstance(net, torch.nn.ReLU) or isinstance(net, torch.nn.PReLU) or isinstance(net,torch.nn.Sigmoid) or isinstance(net, HSwish) or isinstance(net, Swish):
521 | handles.append(net.register_forward_hook(relu_hook))
522 | elif isinstance(net, torch.nn.MaxPool2d) or isinstance(net, torch.nn.AvgPool2d):
523 | handles.append(net.register_forward_hook(pooling_hook))
524 | elif isinstance(net, torch.nn.AdaptiveAvgPool2d):
525 | handles.append(net.register_forward_hook(pooling_hook_ad))
526 | else:
527 | print("warning" + str(net))
528 | return
529 | for c in childrens:
530 | foo(c)
531 |
532 | model.eval()
533 | foo(model)
534 | input = Variable(torch.rand(3, input_res[1], input_res[0]).unsqueeze(0), requires_grad=True)
535 | out = model(input)
536 | total_flops = (sum(list_conv) + sum(list_linear) + sum(list_bn) + sum(list_relu) + sum(list_pooling))
537 | for h in handles:
538 | h.remove()
539 | model.train()
540 | return flops_to_string(total_flops)
541 |
542 |
--------------------------------------------------------------------------------
/feature_extraction/extract_IJB.py:
--------------------------------------------------------------------------------
1 | import os
2 | import cv2
3 | import numpy as np
4 | import shutil
5 | from skimage import transform as trans
6 | from tqdm import tqdm
7 |
8 |
9 | # change this for other dataset
10 | path = "/data/fboutros/IJB_release/IJB_release/IJBB"
11 | image_size = (112,112)
12 | outpath = "/data/maklemt/quality_data"
13 |
14 | ref_lmk = np.array([
15 | [30.2946, 51.6963],
16 | [65.5318, 51.5014],
17 | [48.0252, 71.7366],
18 | [33.5493, 92.3655],
19 | [62.7299, 92.2041]], dtype=np.float32)
20 | ref_lmk[:, 0] += 8.0
21 |
22 | dataset_name = path.split("/")[-1]
23 | rel_img_path = os.path.join(outpath.split("/")[-1], dataset_name, "images")
24 | outpath = os.path.join(outpath, dataset_name)
25 | if not os.path.exists(outpath):
26 | os.makedirs(outpath)
27 | os.makedirs(os.path.join(outpath, "images"))
28 |
29 | print("extract:", dataset_name)
30 |
31 | img_path = os.path.join(path, "loose_crop")
32 | img_list_path = os.path.join(path, "meta", f"{dataset_name.lower()}_name_5pts_score.txt")
33 | img_list = open(img_list_path)
34 | files_list = img_list.readlines()
35 |
36 | txt_file = open(os.path.join(outpath, "image_path_list.txt"), "w")
37 |
38 | for img_index, each_line in tqdm(enumerate(files_list), total=len(files_list)):
39 | name_lmk_score = each_line.strip().split(' ')
40 | img_name = os.path.join(img_path, name_lmk_score[0])
41 | img = cv2.imread(img_name)
42 | lmk = np.array([float(x) for x in name_lmk_score[1:-1]],
43 | dtype=np.float32)
44 | lmk = lmk.reshape((5, 2))
45 |
46 | assert lmk.shape[0] == 5 and lmk.shape[1] == 2
47 |
48 | tform = trans.SimilarityTransform()
49 | tform.estimate(lmk, ref_lmk)
50 | M = tform.params[0:2, :]
51 | img = cv2.warpAffine(img, M, image_size, borderValue=0.0)
52 | img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
53 |
54 | cv2.imwrite(os.path.join(outpath, "images", name_lmk_score[0]), img)
55 | txt_file.write(os.path.join(rel_img_path, name_lmk_score[0])+"\n")
56 |
57 |
58 | txt_file.close()
59 | shutil.copy(
60 | os.path.join(path, "meta", f"{dataset_name.lower()}_template_pair_label.txt"),
61 | os.path.join(outpath, "pair_list.txt")
62 | )
63 | print("pair_list saved")
64 |
--------------------------------------------------------------------------------
/feature_extraction/extract_adience.py:
--------------------------------------------------------------------------------
1 | import os
2 | import cv2
3 | import numpy as np
4 | from tqdm import tqdm
5 | from align_trans import norm_crop
6 |
7 | from mtcnn import MTCNN
8 | import tensorflow as tf
9 |
10 | config = tf.compat.v1.ConfigProto(gpu_options =
11 | tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=0.8)#, device_count = {'GPU': 0}
12 | )
13 | config.gpu_options.allow_growth = True
14 | config.gpu_options.visible_device_list= '0'
15 | session = tf.compat.v1.Session(config=config)
16 | tf.compat.v1.keras.backend.set_session(session)
17 |
18 |
19 | imgs_path = "/data/maklemt/adience"
20 | outpath = "/data/maklemt/quality_data"
21 |
22 |
23 | def find_central_face(img, keypoints):
24 | # if multiple faces are detected, select the face most in the middle of the image
25 | mid_face_idx = 0
26 | if len(keypoints) > 1:
27 | img_mid_point = np.array([img.shape[1]//2, img.shape[0]//2]) # [x (width), y (height)]
28 | noses = np.array([keypoint['keypoints']['nose'] for keypoint in keypoints])
29 | distances = np.linalg.norm(noses - img_mid_point, axis=1) # calculate distance between nose and img mid point
30 | mid_face_idx = np.argmin(distances)
31 |
32 | facial5points = [keypoints[mid_face_idx]['keypoints']['left_eye'], keypoints[mid_face_idx]['keypoints']['right_eye'],
33 | keypoints[mid_face_idx]['keypoints']['nose'], keypoints[mid_face_idx]['keypoints']['mouth_left'],
34 | keypoints[mid_face_idx]['keypoints']['mouth_right']]
35 | return np.array(facial5points)
36 |
37 |
38 |
39 | detector = MTCNN(min_face_size=20, steps_threshold=[0.6, 0.7, 0.9], scale_factor=0.85)
40 | skipped_imgs = []
41 |
42 | dataset_name = imgs_path.split("/")[-1]
43 | rel_img_path = os.path.join(outpath.split("/")[-1], dataset_name, "images")
44 | outpath = os.path.join(outpath, dataset_name)
45 | if not os.path.exists(outpath):
46 | os.makedirs(outpath)
47 | os.makedirs(os.path.join(outpath, "images"))
48 |
49 | print("extract:", dataset_name)
50 |
51 | txt_file = open(os.path.join(outpath, "image_path_list.txt"), "w")
52 | img_files = os.listdir(imgs_path)
53 |
54 | for img_index, img_file in tqdm(enumerate(img_files), total=len(img_files)):
55 | if img_file.split('.')[-1] != "jpg":
56 | continue
57 | img_path = os.path.join(imgs_path, img_file)
58 | img = cv2.imread(img_path)
59 |
60 | keypoints = detector.detect_faces(img)
61 | if len(keypoints) == 0:
62 | skipped_imgs.append(img_file)
63 | continue
64 | facial5points = find_central_face(img, keypoints)
65 | warped_face = norm_crop(img, landmark=facial5points, createEvalDB=True)
66 |
67 | img = cv2.cvtColor(warped_face, cv2.COLOR_RGB2BGR)
68 | person = img_file.split('.')[1]
69 | new_img_name = f"p_{person}_img_{img_index}.jpg"
70 |
71 | cv2.imwrite(os.path.join(outpath, "images", new_img_name), img)
72 | txt_file.write(os.path.join(rel_img_path, new_img_name)+"\n")
73 |
74 | txt_file.close()
75 |
76 |
77 | print("creating pair list...")
78 | pair_list = open(os.path.join(outpath, "pair_list.txt"), "w")
79 | aligned_img_path = os.listdir(os.path.join(outpath, "images"))
80 |
81 | for img_index1, img_file1 in enumerate(aligned_img_path):
82 | if img_file1.split('.')[-1] != "jpg":
83 | continue
84 | person1 = img_file1.split('_')[1]
85 |
86 | for img_index2, img_file2 in enumerate(aligned_img_path[img_index1+1:]):
87 | if img_file2.split('.')[-1] != "jpg":
88 | continue
89 | person2 = img_file2.split('_')[1]
90 | genuine = person1 == person2
91 | pair_list.write(f"{img_file1.split('.')[0]} {img_file2.split('.')[0]} {int(genuine)}\n")
92 |
93 | pair_list.close()
94 | print("pair_list saved")
95 | print("No faces detected in:")
96 | print(skipped_imgs)
97 | print("Total amount of images with no detected face: ", len(skipped_imgs))
98 |
--------------------------------------------------------------------------------
/feature_extraction/extract_bin.py:
--------------------------------------------------------------------------------
1 | import os
2 | import pickle
3 | import cv2
4 | import numpy as np
5 | import mxnet as mx
6 | from tqdm import tqdm
7 |
8 |
9 | # change this for other dataset
10 | path = "data/lfw.bin"
11 | image_size = (112,112)
12 | outpath = "data/quality_data"
13 |
14 |
15 | try:
16 | with open(path, 'rb') as f:
17 | bins, issame_list = pickle.load(f) # py2
18 | except UnicodeDecodeError as e:
19 | with open(path, 'rb') as f:
20 | bins, issame_list = pickle.load(f, encoding='bytes') # py3
21 |
22 | dataset_name = path.split("/")[-1].split(".")[0]
23 | rel_img_path = os.path.join(outpath.split("/")[-1], dataset_name, "images")
24 | outpath = os.path.join(outpath, dataset_name)
25 | if not os.path.exists(outpath):
26 | os.makedirs(outpath)
27 | os.makedirs(os.path.join(outpath, "images"))
28 |
29 | print("extract:", dataset_name)
30 | pair_list = np.zeros((len(issame_list), 3), np.int16)
31 |
32 | txt_file = open(os.path.join(outpath, "image_path_list.txt"), "w")
33 |
34 | for idx in tqdm(range(len(bins))):
35 | _bin = bins[idx]
36 | img = mx.image.imdecode(_bin)
37 | if img.shape[1] != image_size[0]:
38 | img = mx.image.resize_short(img, image_size[0])
39 |
40 | cv2.imwrite(os.path.join(outpath, "images", "%05d.jpg" % idx), img.asnumpy())
41 | txt_file.write(os.path.join(rel_img_path, "%05d.jpg" % idx)+"\n")
42 |
43 | if idx % 2 == 0:
44 | pair_list[idx//2] = [idx, idx+1, issame_list[idx//2]]
45 |
46 | txt_file.close()
47 | np.savetxt(os.path.join(outpath, "pair_list.txt"), pair_list, fmt='%05d')
48 | print("pair_list saved")
49 |
--------------------------------------------------------------------------------
/feature_extraction/extract_emb.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import os
3 | import sys
4 |
5 | import numpy as np
6 | from tqdm import tqdm
7 |
8 |
9 | def parse_arguments(argv):
10 | parser = argparse.ArgumentParser()
11 |
12 | parser.add_argument('--dataset_path', type=str, default='./data/quality_data/XQLFW',
13 | help='dataset directory')
14 | parser.add_argument('--modelname', type=str, default='ElasticFaceModel',
15 | help='ArcFaceModel, CurricularFaceModel, ElasticFaceModel, MagFaceModel')
16 | parser.add_argument('--gpu_id', type=int, default=0,
17 | help='GPU id.')
18 | parser.add_argument('--model_path', type=str, default="",
19 | help='path to pretrained model.')
20 | parser.add_argument('--model_id', type=str, default="295672",
21 | help='digit number in backbone file name')
22 | parser.add_argument('--relative_dir', type=str, default='./data',
23 | help='path to save the embeddings')
24 | parser.add_argument('--color_channel', type=str, default="BGR",
25 | help='input image color channel, two option RGB or BGR')
26 | return parser.parse_args(argv)
27 |
28 |
29 |
30 | def read_img_path_list(image_path_file, relative_dir):
31 | with open(image_path_file, "r") as f:
32 | lines = f.readlines()
33 | lines = [os.path.join(relative_dir, line.rstrip()) for line in lines]
34 | return lines
35 |
36 | def main(param):
37 | dataset_path = param.dataset_path
38 | modelname = param.modelname
39 | gpu_id = param.gpu_id
40 | data_path = param.relative_dir
41 | dataset_name = dataset_path.split("/")[-1]
42 | out_path = os.path.join(data_path, "quality_embeddings", f"{dataset_name}_{modelname}")
43 | if not os.path.exists(out_path):
44 | os.makedirs(out_path)
45 | image_path_list = read_img_path_list(os.path.join(dataset_path, "image_path_list.txt"), data_path)
46 | if modelname == "ArcFaceModel":
47 | from model.ArcFaceModel import ArcFaceModel
48 | face_model = ArcFaceModel(param.model_path, param.model_id, gpu_id=gpu_id)
49 | elif modelname == "CurricularFaceModel":
50 | from model.CurricularFaceModel import CurricularFaceModel
51 | face_model = CurricularFaceModel(param.model_path, param.model_id, gpu_id)
52 | elif modelname == "ElasticFaceModel":
53 | from model.ElasticFaceModel import ElasticFaceModel
54 | face_model = ElasticFaceModel(param.model_path, param.model_id, gpu_id)
55 | elif modelname == "MagFaceModel":
56 | from model.MagFaceModel import MagFaceModel
57 | face_model = MagFaceModel(param.model_path, param.model_id, gpu_id)
58 | else:
59 | print("Unknown model")
60 | exit()
61 |
62 | features = face_model.get_batch_feature(image_path_list)
63 | features_flipped = face_model.get_batch_feature(image_path_list, flip=1,color=param.color_channel)
64 |
65 | # too slow for IJBC
66 | # conc_features = np.concatenate((features, features_flipped), axis=1)
67 | # print(conc_features.shape)
68 |
69 | print("save features")
70 | for i in tqdm(range(len(features))):
71 | conc_features = np.concatenate((features[i], features_flipped[i]), axis=0)
72 | filename = str(str(image_path_list[i]).split("/")[-1].split(".")[0])
73 | np.save(os.path.join(out_path, filename), conc_features)
74 |
75 | if __name__ == '__main__':
76 | main(parse_arguments(sys.argv[1:]))
77 |
--------------------------------------------------------------------------------
/feature_extraction/extract_xqlfw.py:
--------------------------------------------------------------------------------
1 | import os
2 | import cv2
3 |
4 | path = "./data/XQLFW"
5 | outpath = "./data/quality_data"
6 |
7 | dataset_name = path.split("/")[-1]
8 | rel_img_path = os.path.join(outpath.split("/")[-1], dataset_name, "images")
9 | outpath = os.path.join(outpath, dataset_name)
10 | if not os.path.exists(outpath):
11 | os.makedirs(outpath)
12 | os.makedirs(os.path.join(outpath, "images"))
13 |
14 | align_path = os.path.join(path, "xqlfw_aligned_112")
15 |
16 |
17 | def copy_img(person, img_id):
18 | img_name = f"{person}_{str(img_id).zfill(4)}.jpg"
19 | tmp_path = os.path.join(align_path, person, img_name)
20 | img = cv2.imread(tmp_path)
21 | img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
22 | cv2.imwrite(os.path.join(outpath, "images", img_name), img)
23 | return img_name
24 |
25 |
26 | def create_xqlfw_pairs(pairs_filename):
27 | """ reads xqlfw pairs.txt and creates pair_list.txt
28 | and image_path_list.txt of required format
29 | :param pairs_filename: path to pairs.txt
30 | """
31 | txt_file = open(os.path.join(outpath, "image_path_list.txt"), "w")
32 | pair_list = open(os.path.join(outpath, "pair_list.txt"), "w")
33 |
34 | f = open(pairs_filename, 'r')
35 | for line in f.readlines()[1:]:
36 | pair = line.strip().split()
37 | if len(pair) == 3:
38 | img_name1 = copy_img(pair[0], pair[1])
39 | img_name2 = copy_img(pair[0], pair[2])
40 | else:
41 | img_name1 = copy_img(pair[0], pair[1])
42 | img_name2 = copy_img(pair[2], pair[3])
43 |
44 | txt_file.write(os.path.join(rel_img_path, img_name1)+"\n")
45 | txt_file.write(os.path.join(rel_img_path, img_name2)+"\n")
46 | pair_list.write(f"{img_name1} {img_name2} {int(len(pair)==3)}\n")
47 |
48 |
49 | f.close()
50 | txt_file.close()
51 | pair_list.close()
52 |
53 | create_xqlfw_pairs("/data/maklemt/XQLFW/xqlfw_pairs.txt")
54 | print("XQLFW successfully extracted")
--------------------------------------------------------------------------------
/feature_extraction/ijb_pair_file.py:
--------------------------------------------------------------------------------
1 | import os
2 | import pandas as pd
3 | import numpy as np
4 | import sklearn.preprocessing
5 | from tqdm import tqdm
6 | from multiprocessing import Process, Queue
7 |
8 | dataset_name = "IJBC"
9 | model = "ArcFaceModel" # ArcFaceModel, CurricularFaceModel, ElasticFaceModel, MagFaceModel
10 | quality_model = "rankIQA" # Serfiq, BRISQUE, SDD, rankIQ, PFE, magface, FaceQnet, DeepIQA, CRFIQAL
11 | threads = 20
12 | quality_score_path = os.path.join("quality_scores", f"{quality_model}_{dataset_name}.txt")
13 | path = os.path.join("/data/fboutros/IJB_release/IJB_release", dataset_name)
14 | raw_feature_path = os.path.join("/data/maklemt/quality_embeddings", dataset_name + "_" + model + "_raw")
15 | outpath = raw_feature_path[:-4]
16 | if not os.path.exists(outpath):
17 | os.makedirs(outpath)
18 |
19 | tid_mid_path = os.path.join(path, "meta", f"{dataset_name.lower()}_face_tid_mid.txt")
20 | pair_path = os.path.join(path, "meta", f"{dataset_name.lower()}_template_pair_label.txt")
21 |
22 | def read_template_pair_list(path):
23 | pairs = pd.read_csv(path, sep=' ', header=None).values
24 | t1 = pairs[:, 0].astype(np.int)
25 | t2 = pairs[:, 1].astype(np.int)
26 | label = pairs[:, 2].astype(np.int)
27 | print("Template pair list loaded")
28 | return t1, t2, label
29 |
30 |
31 | def read_template_media_list(path):
32 | ijb_meta = pd.read_csv(path, sep=' ', header=None).values
33 | img_names = ijb_meta[:, 0]
34 | templates = ijb_meta[:, 1].astype(np.int)
35 | medias = ijb_meta[:, 2].astype(np.int)
36 | print("Tid mid list loaded")
37 | return img_names, templates, medias
38 |
39 |
40 | def load_raw_templates(files):
41 | features = np.zeros((len(files), 1024))
42 | for i, file in enumerate(files):
43 | filepath = os.path.join(raw_feature_path, file.replace(".jpg", ".npy"))
44 | features[i] = np.load(filepath)
45 | return features
46 |
47 |
48 | def aggregate_templates():
49 | imgs_names, templates, medias = read_template_media_list(tid_mid_path)
50 |
51 | unique_templates = np.unique(templates)
52 |
53 | for uqt in tqdm(unique_templates, total=len(unique_templates)):
54 |
55 | (ind_t,) = np.where(templates == uqt)
56 | face_norm_feats = load_raw_templates(imgs_names[ind_t])
57 | face_medias = medias[ind_t]
58 | unique_medias, unique_media_counts = np.unique(face_medias,
59 | return_counts=True)
60 | media_norm_feats = []
61 | for u, ct in zip(unique_medias, unique_media_counts):
62 | (ind_m,) = np.where(face_medias == u)
63 | if ct == 1:
64 | media_norm_feats += [face_norm_feats[ind_m]]
65 | else: # image features from the same video will be aggregated into one feature
66 | media_norm_feats += [
67 | np.mean(face_norm_feats[ind_m], axis=0, keepdims=True)
68 | ]
69 | media_norm_feats = np.array(media_norm_feats, dtype=np.float32)
70 | aggregated_feature = sklearn.preprocessing.normalize(np.sum(media_norm_feats, axis=0))
71 | np.save(os.path.join(outpath, f"{uqt}.npy"), aggregated_feature[0])
72 |
73 |
74 | aggregate_templates()
75 |
76 |
77 |
78 | def load_quality_scores(path):
79 | quality_scores = pd.read_csv(path, sep=' ', header=None).values
80 | score_dict = { qs[0].split("/")[-1] : float(qs[1]) for qs in quality_scores }
81 | print("Quality scores loaded")
82 | return score_dict
83 |
84 | imgs_names, templates, _ = read_template_media_list(tid_mid_path)
85 | score_dict = load_quality_scores(quality_score_path)
86 |
87 | def get_min_score(person):
88 | (ind_t1,) = np.where(templates == person)
89 | scores = []
90 | for img_name in imgs_names[ind_t1]:
91 | scores.append(score_dict[img_name])
92 | return max(scores) if quality_model == "PFE" else min(scores)
93 |
94 |
95 | def get_score_part(t1s, t2s, labels, queue):
96 | q_pair_list = ""
97 | for t1, t2, label in tqdm(zip(t1s, t2s, labels), total=len(t1s)):
98 | min_score_t1 = get_min_score(t1)
99 | min_score_t2 = get_min_score(t2)
100 | if quality_model == "PFE":
101 | min_score = max([min_score_t1, min_score_t2])
102 | else:
103 | min_score = min([min_score_t1, min_score_t2])
104 |
105 | q_pair_list += f"{t1} {t2} {label} {min_score}\n"
106 |
107 | queue.put(q_pair_list)
108 |
109 |
110 | def create_quality_list():
111 | t1s, t2s, labels = read_template_pair_list(pair_path)
112 | part_idx = len(t1s) // threads
113 | print(f"Quality estimation model: {quality_model}\n{threads+1} Threads")
114 |
115 | q = Queue()
116 | processes = []
117 |
118 | for idx in range(threads):
119 | t1_part = t1s[idx*part_idx:(idx+1)*part_idx]
120 | t2_part = t2s[idx*part_idx:(idx+1)*part_idx]
121 | label_part = labels[idx*part_idx:(idx+1)*part_idx]
122 | p = Process(target=get_score_part, args=(t1_part, t2_part, label_part, q))
123 | processes.append(p)
124 | p.start()
125 |
126 | t1_part = t1s[threads*part_idx:]
127 | t2_part = t2s[threads*part_idx:]
128 | label_part = labels[threads*part_idx:]
129 | p = Process(target=get_score_part, args=(t1_part, t2_part, label_part, q))
130 | processes.append(p)
131 | p.start()
132 |
133 |
134 | save_path = os.path.join("/data/maklemt/quality_data", dataset_name)
135 | pair_list = open(os.path.join(save_path, f"quality_pair_list_{quality_model}_{dataset_name}.txt"), "w")
136 |
137 | for p in processes:
138 | ret = q.get() # will block
139 | pair_list.write(ret)
140 | for p in processes:
141 | p.join()
142 |
143 | pair_list.close()
144 |
145 |
146 | create_quality_list()
--------------------------------------------------------------------------------
/feature_extraction/model/ArcFaceModel.py:
--------------------------------------------------------------------------------
1 | import mxnet as mx
2 |
3 | from model.FaceModel import FaceModel
4 |
5 |
6 | class ArcFaceModel(FaceModel):
7 | def __init__(self, model_prefix, model_epoch, gpu_id):
8 | super(ArcFaceModel, self).__init__(model_prefix, model_epoch, gpu_id)
9 | def _get_model(self, ctx, image_size, prefix, epoch, layer):
10 | print('loading', prefix, epoch)
11 | if ctx>=0:
12 | ctx = mx.gpu(ctx)
13 | else:
14 | ctx = mx.cpu()
15 | sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, int(epoch))
16 | all_layers = sym.get_internals()
17 | sym = all_layers[layer + '_output']
18 | model = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
19 | model.bind(data_shapes=[('data', (1, 3, image_size[0], image_size[1]))])
20 | model.set_params(arg_params, aux_params)
21 | return model
22 |
23 | def _getFeatureBlob(self,input_blob):
24 | data = mx.nd.array(input_blob)
25 | db = mx.io.DataBatch(data=(data,))
26 | self.model.forward(db, is_train=False)
27 | emb = self.model.get_outputs()[0].asnumpy()
28 | return emb
--------------------------------------------------------------------------------
/feature_extraction/model/CurricularFaceModel.py:
--------------------------------------------------------------------------------
1 | import os
2 |
3 | from backbones.model_irse import IR_SE_101, IR_101
4 |
5 | import torch
6 |
7 | from model.FaceModel import FaceModel
8 |
9 |
10 | class CurricularFaceModel(FaceModel):
11 | def __init__(self, model_prefix, model_epoch, gpu_id):
12 | super(CurricularFaceModel, self).__init__(model_prefix, model_epoch, gpu_id)
13 |
14 | def _get_model(self, ctx, image_size, prefix, epoch, layer):
15 | weight = torch.load(os.path.join(prefix,"CurricularFace_Backbone.pth"))
16 | backbone = IR_101([112,112]).to(f"cuda:{ctx}")
17 | backbone.load_state_dict(weight)
18 | model = torch.nn.DataParallel(backbone, device_ids=[ctx])
19 | model.eval()
20 | return model
21 |
22 | @torch.no_grad()
23 | def _getFeatureBlob(self,input_blob):
24 | imgs = torch.Tensor(input_blob).cuda()
25 | imgs.div_(255).sub_(0.5).div_(0.5)
26 | feat,_ = self.model(imgs)
27 | #feat = feat.reshape([self.batch_size, 2 * feat.shape[1]])
28 | return feat.cpu().numpy()
29 |
--------------------------------------------------------------------------------
/feature_extraction/model/ElasticFaceModel.py:
--------------------------------------------------------------------------------
1 | import os
2 |
3 | from backbones.iresnet import iresnet100
4 | from model.FaceModel import FaceModel
5 | import torch
6 | class ElasticFaceModel(FaceModel):
7 | def __init__(self, model_prefix, model_epoch, gpu_id):
8 | super(ElasticFaceModel, self).__init__(model_prefix, model_epoch, gpu_id)
9 |
10 | def _get_model(self, ctx, image_size, prefix, epoch, layer):
11 | weight = torch.load(os.path.join(prefix,epoch+"backbone.pth"))
12 | backbone = iresnet100().to(f"cuda:{ctx}")
13 | backbone.load_state_dict(weight)
14 | model = torch.nn.DataParallel(backbone, device_ids=[ctx])
15 | model.eval()
16 | return model
17 |
18 | @torch.no_grad()
19 | def _getFeatureBlob(self,input_blob):
20 | imgs = torch.Tensor(input_blob).cuda()
21 | imgs.div_(255).sub_(0.5).div_(0.5)
22 | feat = self.model(imgs)
23 | #feat = feat.reshape([self.batch_size, 2 * feat.shape[1]])
24 | return feat.cpu().numpy()
25 |
--------------------------------------------------------------------------------
/feature_extraction/model/FaceModel.py:
--------------------------------------------------------------------------------
1 | import cv2
2 | import numpy as np
3 | from sklearn.preprocessing import normalize
4 |
5 |
6 | class FaceModel():
7 | def __init__(self,model_prefix, model_epoch, ctx_id=7 ):
8 | self.gpu_id=ctx_id
9 | self.image_size = (112, 112)
10 | self.model_prefix=model_prefix
11 | self.model_epoch=model_epoch
12 | self.model=self._get_model(ctx=ctx_id,image_size=self.image_size,prefix=self.model_prefix,epoch=self.model_epoch,layer='fc1')
13 | def _get_model(self, ctx, image_size, prefix, epoch, layer):
14 | pass
15 |
16 | def _getFeatureBlob(self,input_blob):
17 | pass
18 |
19 | def get_feature(self, image_path):
20 | image = cv2.imread(image_path)
21 | image = cv2.resize(image, (112, 112))
22 | a = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
23 | a = np.transpose(a, (2, 0, 1))
24 | input_blob = np.expand_dims(a, axis=0)
25 | emb=self._getFeatureBlob(input_blob)
26 | emb = normalize(emb.reshape(1, -1))
27 | #print(type(emb), emb.shape)
28 | return emb
29 |
30 | def get_batch_feature(self, image_path_list, batch_size=64, flip=0, color="BGR"):
31 | count = 0
32 | num_batch = int(len(image_path_list) / batch_size)
33 | features = []
34 | for i in range(0, len(image_path_list), batch_size):
35 |
36 | if count < num_batch:
37 | tmp_list = image_path_list[i : i+batch_size]
38 | else:
39 | tmp_list = image_path_list[i :]
40 | count += 1
41 |
42 | images = []
43 | for image_path in tmp_list:
44 | image = cv2.imread(image_path)
45 | image = cv2.resize(image, (112, 112))
46 | if (color == "RGB"):
47 | image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
48 | if flip == 1:
49 | image = cv2.flip(image, 1) # flip image horizontally
50 | # a = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
51 | a = np.transpose(image, (2, 0, 1))
52 | images.append(a)
53 | input_blob = np.array(images)
54 |
55 | emb = self._getFeatureBlob(input_blob)
56 | features.append(emb)
57 | print("batch"+str(i))
58 | features = np.vstack(features)
59 | features = normalize(features)
60 | #print('output features:', features.shape)
61 | return features
62 |
--------------------------------------------------------------------------------
/feature_extraction/model/MagFaceModel.py:
--------------------------------------------------------------------------------
1 | import os
2 | from backbones.mag_network_inf import builder_inf
3 | from model.FaceModel import FaceModel
4 | import torch
5 |
6 | class MagFaceModel(FaceModel):
7 | def __init__(self, model_prefix, model_epoch, gpu_id):
8 | super(MagFaceModel, self).__init__(model_prefix, model_epoch, gpu_id)
9 |
10 | def _get_model(self, ctx, image_size, prefix, epoch, layer):
11 | backbone = builder_inf(os.path.join(prefix,"magface_epoch_"+epoch+".pth"), "iresnet100", 512)
12 | model = torch.nn.DataParallel(backbone, device_ids=[ctx])
13 | model.eval()
14 | return model
15 |
16 | @torch.no_grad()
17 | def _getFeatureBlob(self,input_blob):
18 | imgs = torch.Tensor(input_blob).cuda()
19 | imgs.div_(255)
20 | feat = self.model(imgs)
21 | #feat = feat.reshape([self.batch_size, 2 * feat.shape[1]])
22 | return feat.cpu().numpy()
23 |
--------------------------------------------------------------------------------
/losses.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch import nn
3 |
4 | import math
5 | import numpy as np
6 |
7 | def l2_norm(input, axis = 1):
8 | norm = torch.norm(input, 2, axis, True)
9 | output = torch.div(input, norm)
10 |
11 | return output
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 | class CR_FIQA_LOSS(nn.Module):
21 | r"""Implement of ArcFace:
22 | Args:
23 | in_features: size of each input sample
24 | out_features: size of each output sample
25 | s: norm of input feature
26 | m: margin
27 | cos(theta+m)
28 | """
29 |
30 | def __init__(self, in_features, out_features, s=64.0, m=0.50):
31 | super(CR_FIQA_LOSS, self).__init__()
32 | self.in_features = in_features
33 | self.out_features = out_features
34 | self.s = s
35 | self.m = m
36 | self.kernel = nn.Parameter(torch.FloatTensor(in_features, out_features))
37 | nn.init.normal_(self.kernel, std=0.01)
38 |
39 | def forward(self, embbedings, label):
40 | embbedings = l2_norm(embbedings, axis=1)
41 | kernel_norm = l2_norm(self.kernel, axis=0)
42 | cos_theta = torch.mm(embbedings, kernel_norm)
43 | cos_theta = cos_theta.clamp(-1, 1) # for numerical stability
44 | index = torch.where(label != -1)[0]
45 | m_hot = torch.zeros(index.size()[0], cos_theta.size()[1], device=cos_theta.device)
46 | m_hot.scatter_(1, label[index, None], self.m)
47 | with torch.no_grad():
48 | distmat=cos_theta[index,label.view(-1)].detach().clone()
49 | max_negative_cloned=cos_theta.detach().clone()
50 | max_negative_cloned[index,label.view(-1)]= -1e-12
51 | max_negative, _=max_negative_cloned.max(dim=1)
52 | cos_theta.acos_()
53 | cos_theta[index] += m_hot
54 | cos_theta.cos_().mul_(self.s)
55 | return cos_theta, 0 ,distmat[index,None],max_negative[index,None]
--------------------------------------------------------------------------------
/requirement.txt:
--------------------------------------------------------------------------------
1 | tensorboard
2 | easydict
3 | sklearn
4 | matplotlib
5 | pandas
6 | scikit-image
7 | menpo
8 | prettytable
9 | mxnet==1.6.0
10 | pytorch_summary
11 | cv2
12 | torch==1.7.1
--------------------------------------------------------------------------------
/run.sh:
--------------------------------------------------------------------------------
1 | export OMP_NUM_THREADS=4
2 | CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m torch.distributed.launch --nproc_per_node=4 --nnodes=1 \
3 | --node_rank=0 --master_addr="127.0.0.1" --master_port=1236 train_cr_fiqa.py
4 | ps -ef | grep "train" | grep -v grep | awk '{print "kill -9 "$2}' | sh
5 |
--------------------------------------------------------------------------------
/train_cr_fiqa.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import logging
3 | import os
4 | from threading import local
5 | import time
6 |
7 | import torch
8 | import torch.distributed as dist
9 | import torch.nn.functional as F
10 | from torch.nn.parallel.distributed import DistributedDataParallel
11 | import torch.utils.data.distributed
12 | from torch.nn.utils import clip_grad_norm_
13 | from torch.nn import CrossEntropyLoss
14 |
15 | import losses
16 | from config import config as cfg
17 | from dataset import MXFaceDataset, DataLoaderX
18 | from utils.utils_callbacks import CallBackVerification, CallBackLogging, CallBackModelCheckpoint
19 | from utils.utils_logging import AverageMeter, init_logging
20 |
21 | from backbones.iresnet import iresnet100, iresnet50
22 |
23 |
24 | torch.backends.cudnn.benchmark = True
25 |
26 | def main(args):
27 | dist.init_process_group(backend='nccl', init_method='env://')
28 | local_rank = args.local_rank
29 | torch.cuda.set_device(local_rank)
30 | rank = dist.get_rank()
31 | world_size = dist.get_world_size()
32 |
33 | if not os.path.exists(cfg.output) and rank == 0:
34 | os.makedirs(cfg.output)
35 | else:
36 | time.sleep(2)
37 |
38 | log_root = logging.getLogger()
39 | init_logging(log_root, rank, cfg.output)
40 |
41 | trainset = MXFaceDataset(root_dir=cfg.rec, local_rank=local_rank)
42 |
43 | train_sampler = torch.utils.data.distributed.DistributedSampler(
44 | trainset, shuffle=True)
45 |
46 | train_loader = DataLoaderX(
47 | local_rank=local_rank, dataset=trainset, batch_size=cfg.batch_size,
48 | sampler=train_sampler, num_workers=16, pin_memory=True, drop_last=True)
49 |
50 | # load evaluation
51 |
52 | if cfg.network == "iresnet100":
53 | backbone = iresnet100(dropout=0.4,num_features=cfg.embedding_size).to(local_rank)
54 | elif cfg.network == "iresnet50":
55 | backbone = iresnet50(dropout=0.4,num_features=cfg.embedding_size, use_se=False, qs=1).to(local_rank)
56 | else:
57 | backbone = None
58 | logging.info("load backbone failed!")
59 |
60 | if args.resume:
61 | try:
62 | backbone_pth = os.path.join(cfg.output, str(cfg.global_step) + "backbone.pth")
63 | backbone.load_state_dict(torch.load(backbone_pth, map_location=torch.device(local_rank)))
64 | if rank == 0:
65 | logging.info("backbone student loaded successfully!")
66 |
67 | if rank == 0:
68 | logging.info("backbone resume loaded successfully!")
69 | except (FileNotFoundError, KeyError, IndexError, RuntimeError):
70 | logging.info("load backbone resume init, failed!")
71 |
72 | for ps in backbone.parameters():
73 | dist.broadcast(ps, 0)
74 |
75 | backbone = DistributedDataParallel(
76 | module=backbone, broadcast_buffers=False, device_ids=[local_rank])
77 | backbone.train()
78 |
79 | # get header
80 | if args.loss == "CR_FIQA_LOSS":
81 | header = losses.CR_FIQA_LOSS(in_features=cfg.embedding_size, out_features=cfg.num_classes, s=cfg.s, m=cfg.m).to(local_rank)
82 | else:
83 | print("Header not implemented")
84 | if args.resume:
85 | try:
86 | header_pth = os.path.join(cfg.identity_model, str(cfg.identity_step) + "header.pth")
87 | header.load_state_dict(torch.load(header_pth, map_location=torch.device(local_rank)))
88 |
89 | if rank == 0:
90 | logging.info("header resume loaded successfully!")
91 | except (FileNotFoundError, KeyError, IndexError, RuntimeError):
92 | logging.info("header resume init, failed!")
93 |
94 | header = DistributedDataParallel(
95 | module=header, broadcast_buffers=False, device_ids=[local_rank])
96 | header.train()
97 |
98 | opt_backbone = torch.optim.SGD(
99 | params=[{'params': backbone.parameters()}],
100 | lr=cfg.lr / 512 * cfg.batch_size * world_size,
101 | momentum=0.9, weight_decay=cfg.weight_decay)
102 | opt_header = torch.optim.SGD(
103 | params=[{'params': header.parameters()}],
104 | lr=cfg.lr / 512 * cfg.batch_size * world_size,
105 | momentum=0.9, weight_decay=cfg.weight_decay)
106 |
107 | scheduler_backbone = torch.optim.lr_scheduler.LambdaLR(
108 | optimizer=opt_backbone, lr_lambda=cfg.lr_func)
109 | scheduler_header = torch.optim.lr_scheduler.LambdaLR(
110 | optimizer=opt_header, lr_lambda=cfg.lr_func)
111 |
112 | criterion = CrossEntropyLoss()
113 | criterion_qs= torch.nn.SmoothL1Loss(beta=0.5)
114 |
115 | start_epoch = 0
116 | total_step = int(len(trainset) / cfg.batch_size / world_size * cfg.num_epoch)
117 | if rank == 0: logging.info("Total Step is: %d" % total_step)
118 |
119 | if args.resume:
120 | rem_steps = (total_step - cfg.global_step)
121 | cur_epoch = cfg.num_epoch - int(cfg.num_epoch / total_step * rem_steps)
122 | logging.info("resume from estimated epoch {}".format(cur_epoch))
123 | logging.info("remaining steps {}".format(rem_steps))
124 |
125 | start_epoch = cur_epoch
126 | scheduler_backbone.last_epoch = cur_epoch
127 | scheduler_header.last_epoch = cur_epoch
128 |
129 | # --------- this could be solved more elegant ----------------
130 | opt_backbone.param_groups[0]['lr'] = scheduler_backbone.get_lr()[0]
131 | opt_header.param_groups[0]['lr'] = scheduler_header.get_lr()[0]
132 |
133 | print("last learning rate: {}".format(scheduler_header.get_lr()))
134 | # ------------------------------------------------------------
135 |
136 | callback_verification = CallBackVerification(cfg.eval_step, rank, cfg.val_targets, cfg.rec)
137 | callback_logging = CallBackLogging(50, rank, total_step, cfg.batch_size, world_size, writer=None)
138 | callback_checkpoint = CallBackModelCheckpoint(rank, cfg.output)
139 | alpha=10.0 #10.0
140 | loss = AverageMeter()
141 | global_step = cfg.global_step
142 | for epoch in range(start_epoch, cfg.num_epoch):
143 | train_sampler.set_epoch(epoch)
144 | for _, (img, label) in enumerate(train_loader):
145 | global_step += 1
146 | img = img.cuda(local_rank, non_blocking=True)
147 | label = label.cuda(local_rank, non_blocking=True)
148 |
149 | features, qs = backbone(img)
150 | thetas, std, ccs,nnccs = header(features, label)
151 | loss_qs=criterion_qs(ccs/ nnccs,qs)
152 | loss_v = criterion(thetas, label) + alpha* loss_qs
153 | loss_v.backward()
154 | clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)
155 |
156 | opt_backbone.step()
157 | opt_header.step()
158 |
159 | opt_backbone.zero_grad()
160 | opt_header.zero_grad()
161 |
162 | loss.update(loss_v.item(), 1)
163 |
164 | callback_logging(global_step, loss, epoch, 0,loss_qs)
165 | callback_verification(global_step, backbone)
166 |
167 | scheduler_backbone.step()
168 | scheduler_header.step()
169 |
170 | callback_checkpoint(global_step, backbone, header)
171 |
172 | dist.destroy_process_group()
173 |
174 |
175 | if __name__ == "__main__":
176 | parser = argparse.ArgumentParser(description='PyTorch CR-FIQA Training')
177 | parser.add_argument('--local_rank', type=int, default=0, help='local_rank')
178 | parser.add_argument('--loss', type=str, default="CR_FIQA_LOSS", help="loss function")
179 | parser.add_argument('--resume', type=int, default=0, help="resume training")
180 | args_ = parser.parse_args()
181 | main(args_)
182 |
--------------------------------------------------------------------------------
/utils/NIST.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/fdbtrs/CR-FIQA/a8f98da6ee03db1e8dc091b5910931d41d69fa44/utils/NIST.png
--------------------------------------------------------------------------------
/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/fdbtrs/CR-FIQA/a8f98da6ee03db1e8dc091b5910931d41d69fa44/utils/__init__.py
--------------------------------------------------------------------------------
/utils/align_trans.py:
--------------------------------------------------------------------------------
1 | # -*- coding: utf-8 -*-
2 | import numpy as np
3 | import cv2
4 | from skimage import transform as trans
5 |
6 |
7 | arcface_ref_points = np.array([ [30.2946, 51.6963],
8 | [65.5318, 51.5014],
9 | [48.0252, 71.7366],
10 | [33.5493, 92.3655],
11 | [62.7299, 92.2041]
12 | ], dtype=np.float32 )
13 |
14 | # https://github.com/deepinsight/insightface/blob/master/python-package/insightface/utils/face_align.py
15 | # [:,0] += 8.0
16 | arcface_eval_ref_points = np.array([
17 | [38.2946, 51.6963],
18 | [73.5318, 51.5014],
19 | [56.0252, 71.7366],
20 | [41.5493, 92.3655],
21 | [70.7299, 92.2041]
22 | ], dtype=np.float32)
23 |
24 | # lmk is prediction; src is template
25 | def estimate_norm(lmk, image_size=112, createEvalDB=False):
26 | """ estimate the transformation matrix
27 | :param lmk: detected landmarks
28 | :param image_size: resulting image size (default=112)
29 | :param createEvalDB: (boolean) crop an evaluation or training dataset
30 | :return: transformation matrix M and index
31 | """
32 | assert lmk.shape == (5, 2)
33 | assert image_size == 112
34 | tform = trans.SimilarityTransform()
35 | lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)
36 | min_M = []
37 | min_index = []
38 | min_error = float('inf')
39 | if createEvalDB:
40 | src = arcface_eval_ref_points
41 | else:
42 | src = arcface_ref_points
43 | src = np.expand_dims(src, axis=0)
44 |
45 | for i in np.arange(src.shape[0]):
46 | tform.estimate(lmk, src[i])
47 | M = tform.params[0:2, :]
48 | results = np.dot(M, lmk_tran.T)
49 | results = results.T
50 | error = np.sum(np.sqrt(np.sum((results - src[i])**2, axis=1)))
51 | # print(error)
52 | if error < min_error:
53 | min_error = error
54 | min_M = M
55 | min_index = i
56 | return min_M, min_index
57 |
58 |
59 | # norm_crop from Arcface repository (insightface/recognition/common/face_align.py)
60 | def norm_crop(img, landmark, image_size=112, createEvalDB=False):
61 | """ transforms image to match the landmarks with reference landmarks
62 | :param landmark: detected landmarks
63 | :param image_size: resulting image size (default=112)
64 | :param createEvalDB: (boolean) crop an evaluation or training dataset
65 | :return: transformed image
66 | """
67 | M, pose_index = estimate_norm(landmark, image_size=image_size, createEvalDB=createEvalDB)
68 | warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
69 | return warped
70 |
--------------------------------------------------------------------------------
/utils/utils_amp.py:
--------------------------------------------------------------------------------
1 | from typing import Dict, List
2 |
3 | import torch
4 | from torch._six import container_abcs
5 | from torch.cuda.amp import GradScaler
6 |
7 |
8 | class _MultiDeviceReplicator(object):
9 | """
10 | Lazily serves copies of a tensor to requested devices. Copies are cached per-device.
11 | """
12 |
13 | def __init__(self, master_tensor: torch.Tensor) -> None:
14 | assert master_tensor.is_cuda
15 | self.master = master_tensor
16 | self._per_device_tensors: Dict[torch.device, torch.Tensor] = {}
17 |
18 | def get(self, device) -> torch.Tensor:
19 | retval = self._per_device_tensors.get(device, None)
20 | if retval is None:
21 | retval = self.master.to(device=device, non_blocking=True, copy=True)
22 | self._per_device_tensors[device] = retval
23 | return retval
24 |
25 |
26 | class MaxClipGradScaler(GradScaler):
27 | def __init__(self, init_scale, max_scale: float, growth_interval=100):
28 | GradScaler.__init__(self, init_scale=init_scale, growth_interval=growth_interval)
29 | self.max_scale = max_scale
30 |
31 | def scale_clip(self):
32 | if self.get_scale() == self.max_scale:
33 | self.set_growth_factor(1)
34 | elif self.get_scale() < self.max_scale:
35 | self.set_growth_factor(2)
36 | elif self.get_scale() > self.max_scale:
37 | self._scale.fill_(self.max_scale)
38 | self.set_growth_factor(1)
39 |
40 | def scale(self, outputs):
41 | """
42 | Multiplies ('scales') a tensor or list of tensors by the scale factor.
43 |
44 | Returns scaled outputs. If this instance of :class:`GradScaler` is not enabled, outputs are returned
45 | unmodified.
46 |
47 | Arguments:
48 | outputs (Tensor or iterable of Tensors): Outputs to scale.
49 | """
50 | if not self._enabled:
51 | return outputs
52 | self.scale_clip()
53 | # Short-circuit for the common case.
54 | if isinstance(outputs, torch.Tensor):
55 | assert outputs.is_cuda
56 | if self._scale is None:
57 | self._lazy_init_scale_growth_tracker(outputs.device)
58 | assert self._scale is not None
59 | return outputs * self._scale.to(device=outputs.device, non_blocking=True)
60 |
61 | # Invoke the more complex machinery only if we're treating multiple outputs.
62 | stash: List[_MultiDeviceReplicator] = [] # holds a reference that can be overwritten by apply_scale
63 |
64 | def apply_scale(val):
65 | if isinstance(val, torch.Tensor):
66 | assert val.is_cuda
67 | if len(stash) == 0:
68 | if self._scale is None:
69 | self._lazy_init_scale_growth_tracker(val.device)
70 | assert self._scale is not None
71 | stash.append(_MultiDeviceReplicator(self._scale))
72 | return val * stash[0].get(val.device)
73 | elif isinstance(val, container_abcs.Iterable):
74 | iterable = map(apply_scale, val)
75 | if isinstance(val, list) or isinstance(val, tuple):
76 | return type(val)(iterable)
77 | else:
78 | return iterable
79 | else:
80 | raise ValueError("outputs must be a Tensor or an iterable of Tensors")
81 | return apply_scale(outputs)
82 |
--------------------------------------------------------------------------------
/utils/utils_callbacks.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import os
3 | import time
4 | from typing import List
5 |
6 | import torch
7 |
8 | from eval import verification
9 | from utils.utils_logging import AverageMeter
10 |
11 |
12 | class CallBackVerification(object):
13 | def __init__(self, frequent, rank, val_targets, rec_prefix, image_size=(112, 112)):
14 | self.frequent: int = frequent
15 | self.rank: int = rank
16 | self.highest_acc: float = 0.0
17 | self.highest_acc_list: List[float] = [0.0] * len(val_targets)
18 | self.ver_list: List[object] = []
19 | self.ver_name_list: List[str] = []
20 | if self.rank == 0:
21 | self.init_dataset(val_targets=val_targets, data_dir=rec_prefix, image_size=image_size)
22 |
23 | def ver_test(self, backbone: torch.nn.Module, global_step: int):
24 | results = []
25 | for i in range(len(self.ver_list)):
26 | acc1, std1, acc2, std2, xnorm, embeddings_list = verification.test(
27 | self.ver_list[i], backbone, 10, 10)
28 | logging.info('[%s][%d]XNorm: %f' % (self.ver_name_list[i], global_step, xnorm))
29 | logging.info('[%s][%d]Accuracy-Flip: %1.5f+-%1.5f' % (self.ver_name_list[i], global_step, acc2, std2))
30 | if acc2 > self.highest_acc_list[i]:
31 | self.highest_acc_list[i] = acc2
32 | logging.info(
33 | '[%s][%d]Accuracy-Highest: %1.5f' % (self.ver_name_list[i], global_step, self.highest_acc_list[i]))
34 | results.append(acc2)
35 |
36 | def init_dataset(self, val_targets, data_dir, image_size):
37 | for name in val_targets:
38 | path = os.path.join(data_dir, name + ".bin")
39 | if os.path.exists(path):
40 | data_set = verification.load_bin(path, image_size)
41 | self.ver_list.append(data_set)
42 | self.ver_name_list.append(name)
43 |
44 | def __call__(self, num_update, backbone: torch.nn.Module):
45 | if self.rank == 0 and num_update > 0 and num_update % self.frequent == 0:
46 | backbone.eval()
47 | self.ver_test(backbone, num_update)
48 | backbone.train()
49 |
50 |
51 | class CallBackLogging(object):
52 | def __init__(self, frequent, rank, total_step, batch_size, world_size, writer=None, resume=0, rem_total_steps=None):
53 | self.frequent: int = frequent
54 | self.rank: int = rank
55 | self.time_start = time.time()
56 | self.total_step: int = total_step
57 | self.batch_size: int = batch_size
58 | self.world_size: int = world_size
59 | self.writer = writer
60 | self.resume = resume
61 | self.rem_total_steps = rem_total_steps
62 |
63 | self.init = False
64 | self.tic = 0
65 |
66 | def __call__(self, global_step, loss: AverageMeter, epoch: int, std:float, center:float):
67 | if self.rank == 0 and global_step > 0 and global_step % self.frequent == 0:
68 | if self.init:
69 | try:
70 | speed: float = self.frequent * self.batch_size / (time.time() - self.tic)
71 | speed_total = speed * self.world_size
72 | except ZeroDivisionError:
73 | speed_total = float('inf')
74 |
75 | time_now = (time.time() - self.time_start) / 3600
76 | # TODO: resume time_total is not working
77 | if self.resume:
78 | time_total = time_now / ((global_step + 1) / self.rem_total_steps)
79 | else:
80 | time_total = time_now / ((global_step + 1) / self.total_step)
81 | time_for_end = time_total - time_now
82 | if self.writer is not None:
83 | self.writer.add_scalar('time_for_end', time_for_end, global_step)
84 | self.writer.add_scalar('loss', loss.avg, global_step)
85 | msg = "Speed %.2f samples/sec Loss %.4f Margin %.4f Center %.4f Epoch: %d Global Step: %d Required: %1.f hours" % (
86 | speed_total, loss.avg, std, center,epoch, global_step, time_for_end
87 | )
88 | logging.info(msg)
89 | loss.reset()
90 | self.tic = time.time()
91 | else:
92 | self.init = True
93 | self.tic = time.time()
94 |
95 | class CallBackModelCheckpoint(object):
96 | def __init__(self, rank, output="./"):
97 | self.rank: int = rank
98 | self.output: str = output
99 |
100 | def __call__(self, global_step, backbone: torch.nn.Module, header: torch.nn.Module = None):
101 | if global_step > 100 and self.rank == 0:
102 | torch.save(backbone.module.state_dict(), os.path.join(self.output, str(global_step)+ "backbone.pth"))
103 | if global_step > 100 and header is not None:
104 | torch.save(header.module.state_dict(), os.path.join(self.output, str(global_step)+ "header.pth"))
105 |
--------------------------------------------------------------------------------
/utils/utils_logging.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import os
3 | import sys
4 |
5 |
6 | class AverageMeter(object):
7 | """Computes and stores the average and current value
8 | """
9 |
10 | def __init__(self):
11 | self.val = None
12 | self.avg = None
13 | self.sum = None
14 | self.count = None
15 | self.reset()
16 |
17 | def reset(self):
18 | self.val = 0
19 | self.avg = 0
20 | self.sum = 0
21 | self.count = 0
22 |
23 | def update(self, val, n=1):
24 | self.val = val
25 | self.sum += val * n
26 | self.count += n
27 | self.avg = self.sum / self.count
28 |
29 |
30 | def init_logging(log_root, rank, models_root, logfile=None):
31 | if rank is 0:
32 | log_root.setLevel(logging.INFO)
33 | formatter = logging.Formatter("Training: %(asctime)s-%(message)s")
34 | file_name = "training.log" if logfile is None else logfile
35 | handler_file = logging.FileHandler(os.path.join(models_root, file_name))
36 | handler_stream = logging.StreamHandler(sys.stdout)
37 | handler_file.setFormatter(formatter)
38 | handler_stream.setFormatter(formatter)
39 | log_root.addHandler(handler_file)
40 | log_root.addHandler(handler_stream)
41 | log_root.info('rank_id: %d' % rank)
42 |
--------------------------------------------------------------------------------
/utils/workflow.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/fdbtrs/CR-FIQA/a8f98da6ee03db1e8dc091b5910931d41d69fa44/utils/workflow.png
--------------------------------------------------------------------------------