-
Notifications
You must be signed in to change notification settings - Fork 9
/
retrieval_test.py
115 lines (86 loc) · 3.42 KB
/
retrieval_test.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
import pickle
from sklearn.neighbors import NearestNeighbors
import numpy as np
from PIL import Image
import os
import torchvision.transforms as transforms
from utils.visualize import Visualizer
import tqdm
PHOTO_ROOT = '/data1/zzl/dataset/photo-test'
SKETCH_ROOT = '/data1/zzl/dataset/sketch-triplet-test'
photo_data = pickle.load(open('feature/photo-resnet50_64-265.pkl', 'rb'))
sketch_data = pickle.load(open('feature/sketch-resnet50_64-265.pkl', 'rb'))
# print(photo_data['name'][0])
photo_feature = photo_data['feature']
photo_name = photo_data['name']
sketch_feature = sketch_data['feature']
sketch_name = sketch_data['name']
# print(np.size(photo_feature, 0))
nbrs = NearestNeighbors(n_neighbors=np.size(photo_feature, 0),
algorithm='brute', metric='euclidean').fit(photo_feature)
s_len = np.size(sketch_feature, 0)
picked = [0] * s_len
transform = transforms.Compose([
transforms.ToTensor()
])
vis = Visualizer('caffe2torch_test')
count = 0
for ii in range(20):
index = np.random.randint(0, s_len)
while picked[index]:
index = np.random.randint(0, s_len)
picked[index] = 1
query_feature = sketch_feature[index]
query_feature = np.reshape(query_feature, [1, np.shape(query_feature)[0]])
query_name = sketch_name[index]
# print(query_name)
distances, indices = nbrs.kneighbors(query_feature)
query_split = query_name.split('/')
query_class = query_split[0]
query_img = query_split[1]
query_image = np.array(Image.open(os.path.join(SKETCH_ROOT, query_name)).convert('RGB'))
for i, indice in enumerate(indices[0][:5]):
retrievaled_name = photo_name[indice]
retrievaled_im = np.array(Image.open(os.path.join(PHOTO_ROOT, retrievaled_name)).convert('RGB'))
query_image = np.append(query_image, retrievaled_im, axis=1)
retrievaled_class = retrievaled_name.split('/')[0]
retrievaled_name = retrievaled_name.split('/')[1]
retrievaled_name = retrievaled_name.split('.')[0]
if retrievaled_class == query_class:
print(ii, 'correct class', query_name, retrievaled_name)
if query_img.find(retrievaled_name) != -1:
print(ii, 'correct item', query_name)
count += 1
if ii == 0:
result = query_image
else:
result = np.append(result, query_image, axis=0)
result = transform(result)
vis.images(result.numpy(), win='result')
print(count)
count = 0
count_5 = 0
K = 5
div = 0
for ii, (query_sketch, query_name) in tqdm.tqdm(enumerate(zip(sketch_feature, sketch_name))):
query_sketch = np.reshape(query_sketch, [1, np.shape(query_sketch)[0]])
query_split = query_name.split('/')
query_class = query_split[0]
query_img = query_split[1]
distances, indices = nbrs.kneighbors(query_sketch)
div += distances[0][1] - distances[0][0]
# top K
for i, indice in enumerate(indices[0][:K]):
retrievaled_name = photo_name[indice]
retrievaled_class = retrievaled_name.split('/')[0]
retrievaled_name = retrievaled_name.split('/')[1]
retrievaled_name = retrievaled_name.split('.')[0]
if retrievaled_class == query_class:
if query_img.find(retrievaled_name) != -1:
if i == 0:
count += 1
count_5 += 1
break
recall = count / (ii+1)
recall_5 = count_5 / (ii+1)
print('recall@1 :', recall, ' recall@5 :', recall_5, 'div :', div/(ii+1))