facexlib

1
import argparse
2
import glob
3
import math
4
import numpy as np
5
import os
6
import torch
7

8
from facexlib.recognition import ResNetArcFace, cosin_metric, load_image
9

10
if __name__ == '__main__':
11
    parser = argparse.ArgumentParser()
12
    parser.add_argument('--folder1', type=str)
13
    parser.add_argument('--folder2', type=str)
14
    parser.add_argument('--model_path', type=str, default='facexlib/recognition/weights/arcface_resnet18.pth')
15

16
    args = parser.parse_args()
17

18
    img_list1 = sorted(glob.glob(os.path.join(args.folder1, '*')))
19
    img_list2 = sorted(glob.glob(os.path.join(args.folder2, '*')))
20
    print(img_list1, img_list2)
21
    model = ResNetArcFace(block='IRBlock', layers=(2, 2, 2, 2), use_se=False)
22
    model.load_state_dict(torch.load(args.model_path))
23
    model.to(torch.device('cuda'))
24
    model.eval()
25

26
    dist_list = []
27
    identical_count = 0
28
    for idx, (img_path1, img_path2) in enumerate(zip(img_list1, img_list2)):
29
        basename = os.path.splitext(os.path.basename(img_path1))[0]
30
        img1 = load_image(img_path1)
31
        img2 = load_image(img_path2)
32

33
        data = torch.stack([img1, img2], dim=0)
34
        data = data.to(torch.device('cuda'))
35
        output = model(data)
36
        print(output.size())
37
        output = output.data.cpu().numpy()
38
        dist = cosin_metric(output[0], output[1])
39
        dist = np.arccos(dist) / math.pi * 180
40
        print(f'{idx} - {dist} o : {basename}')
41
        if dist < 1:
42
            print(f'{basename} is almost identical to original.')
43
            identical_count += 1
44
        else:
45
            dist_list.append(dist)
46

47
    print(f'Result dist: {sum(dist_list) / len(dist_list):.6f}')
48
    print(f'identical count: {identical_count}')
49