forked from TreB1eN/InsightFace_Pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
inference_dog_eyes_verification.py
156 lines (119 loc) · 6.15 KB
/
inference_dog_eyes_verification.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
import argparse
import cv2
from PIL.JpegImagePlugin import JpegImageFile
cv2.ocl.setUseOpenCL(False)
import albumentations as al
from albumentations.pytorch import ToTensorV2
import numpy as np
import torch
from backbone import Backbone
def get_test_transforms(input_size=None, use_random_crop=False, use_gray=False,
use_center_crop=False,
center_crop_ratio=0.8):
if use_random_crop:
compose = [al.Resize(int(input_size * 1.1), int(input_size * 1.1)),
al.CenterCrop(input_size, input_size)]
elif use_center_crop:
compose = [al.Resize(int(input_size * (2.0 - center_crop_ratio)), int(input_size * (2.0 - center_crop_ratio))),
al.CenterCrop(input_size, input_size)]
else:
compose = [al.Resize(input_size, input_size)]
return al.Compose(compose + [
al.ToGray(p=1.0 if use_gray else 0.0),
al.Normalize(),
ToTensorV2()
])
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
class DogEyesVerifier(object):
def __init__(self, input_size=112, model_path=None, pos_thr=None, neg_thr=None, drop_ratio=0.6,
net_depth=50, use_random_crop=False,
use_gray=False, use_center_crop=False, center_crop_ratio=0.8, device='CPU', use_onnx=False):
assert model_path is not None
self.pos_thr = pos_thr
self.neg_thr = neg_thr
self.input_size = input_size
self.transforms = get_test_transforms(input_size, use_random_crop=use_random_crop, use_gray=use_gray,
use_center_crop=use_center_crop, center_crop_ratio=center_crop_ratio)
self.use_onnx = use_onnx
if use_onnx:
import onnxruntime
self.model = onnxruntime.InferenceSession(model_path)
else:
self.model = Backbone(net_depth, drop_ratio, 'ir_se').to(device)
self.model.eval()
if device == 'cpu':
self.model.load_state_dict(torch.load(model_path, map_location=device), strict=True)
else:
self.model.load_state_dict(torch.load(model_path), strict=True)
self.use_cuda = device == "cuda"
def is_same(self, img1, img2, is_same_side, pos_thr=None, neg_thr=None, use_pos_low_thr=False):
if pos_thr is None:
pos_thr = self.pos_thr
if neg_thr is None:
neg_thr = self.neg_thr
if pos_thr is None and neg_thr is None:
raise ValueError("pos_thr and neg_thr are None.")
if isinstance(img1, str):
img1 = cv2.imread(img1, cv2.IMREAD_COLOR)
img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
elif isinstance(img1, JpegImageFile):
if img1.mode != "RGB":
img1 = img1.convert("RGB")
img1 = np.array(img1)
if isinstance(img2, str):
img2 = cv2.imread(img2, cv2.IMREAD_COLOR)
img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
elif isinstance(img2, JpegImageFile):
if img2.mode != "RGB":
img2 = img2.convert("RGB")
img2 = np.array(img2)
img1 = self.transforms(image=img1)['image'].unsqueeze(0)
img2 = self.transforms(image=img2)['image'].unsqueeze(0)
if self.use_cuda:
img1 = img1.cuda()
img2 = img2.cuda()
if self.use_onnx:
input1 = {self.model.get_inputs()[0].name: to_numpy(img1)}
input2 = {self.model.get_inputs()[0].name: to_numpy(img2)}
embedding1 = self.model.run(None, input1)[0]
embedding2 = self.model.run(None, input2)[0]
else:
with torch.set_grad_enabled(False):
embedding1 = self.model(img1).cpu().data.numpy()
embedding2 = self.model(img2).cpu().data.numpy()
dist = np.sum(np.square(np.subtract(embedding1, embedding2)), 1)[0]
if is_same_side:
return dist < pos_thr
else:
if use_pos_low_thr:
return dist < neg_thr
else:
return False
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='for common image metric learning')
parser.add_argument("-depth", "--net_depth", help="how many layers [50,100,152]", default=50, type=int)
parser.add_argument('--input_size', type=int, default=112)
parser.add_argument('--use_random_crop', default=False, action="store_true")
parser.add_argument('--model_path', default=None, type=str)
parser.add_argument('--use_gpu', default=False, action="store_true")
# onnx 모델 사용할지
parser.add_argument('--use_onnx', default=False, action="store_true")
# 미리 두 눈 이미지를 left vs right 모델로 돌려서 동일한 쪽의 눈인지 체크해줘야 한다.
parser.add_argument('--is_same_side', default=False, action="store_true")
# left vs right 모델이 오탐할 경우에도 두 눈의 vector거리가 가까우면 동일한 눈으로 취급한다.
# 전체적으로 정확도가 더 높아짐. 사용하는것을 추천.
parser.add_argument('--use_pos_low_thr', default=False, action="store_true")
parser.add_argument('--img1', default=None, type=str)
parser.add_argument('--img2', default=None, type=str)
# positive pair(is_same_side가 true) 일때의 threshold값
parser.add_argument('--pos_thr', default=None, type=float)
# negative pair(is_same_side가 false) 인데 use_pos_low_thr 옵션을 사용할 경우 이 negative threshold값을 사용하여 판단
parser.add_argument('--neg_thr', default=None, type=float)
args = parser.parse_args()
verifier = DogEyesVerifier(input_size=args.input_size, model_path=args.model_path,
pos_thr=args.pos_thr, neg_thr=args.neg_thr, use_random_crop=args.use_random_crop,
device="cuda" if args.use_gpu else "cpu", use_onnx=args.use_onnx)
result = verifier.is_same(args.img1, args.img2, args.is_same_side, pos_thr=args.pos_thr, neg_thr=args.neg_thr,
use_pos_low_thr=args.use_pos_low_thr)
print("result", result)