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
|
import random
import torch
from cv.image_similarity.cnn.build_models import IR_50, IR_101, IR_152
import torchvision
import os
from PIL import Image
to_torch_tensor = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize([0.5, 0.5, 0.5],[0.5, 0.5, 0.5])])
def l2_norm(input, axis=1):
norm = torch.norm(input, 2, axis, True)
output = torch.div(input, norm)
return output
def init_model(model, param, device):
m = model([112, 112])
m.eval()
m.to(device)
m.load_state_dict(torch.load(param, map_location=torch.device('cpu')))
return m
def get_model_pool(device):
model_pool = []
# double
model_pool.append(init_model(IR_50, 'models/backbone_ir50_ms1m_epoch120.pth', device))
# model_pool.append(init_model(IR_50, 'models/backbone_ir50_ms1m_epoch120.pth', device))
#
# model_pool.append(init_model(IR_50, 'models/Backbone_IR_50_LFW.pth', device))
# model_pool.append(init_model(IR_101, 'models/Backbone_IR_101_Batch_108320.pth', device))
# model_pool.append(init_model(IR_152, 'models/Backbone_IR_152_MS1M_Epoch_112.pth', device))
return model_pool
def get_model(device):
return init_model(IR_50, 'models/backbone_ir50_ms1m_epoch120.pth', device)
device = torch.device('cpu')
model = get_model(device)
print('----models load over----')
images = os.listdir('../imgs/lena')
vectors = []
for img in images:
print(img)
img = Image.open('../imgs/lena/' + img).convert('RGB')
# print(dir(img))
print(img.size, img.mode)
img = to_torch_tensor(img.resize((112, 112), Image.ANTIALIAS))
img = img.unsqueeze_(0).to(device)
feature = model(img)
vectors.append(l2_norm(feature).detach_())
print('----vectors calculate over----')
for i in range(len(vectors)):
for j in range(len(vectors)):
# consine distance
dist = (vectors[i]*vectors[j]).sum().item()
print(images[i], images[j], dist)
print('-------------')
|
