-
Notifications
You must be signed in to change notification settings - Fork 2
/
cbir.py
144 lines (109 loc) · 3.8 KB
/
cbir.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
import numpy as np
import sys
import matplotlib.pyplot as plt
import imageio as im
import collections as cl
import tqdm as tq
sys.path.append('../extractor')
import feature_extractor_bib as feb
def load_parameters():
k = int(input())
query_img = str(input()).rstrip()
imgs_total = int(input())
dist_f = str(input()).rstrip()
descs_total = int(input())
desc_files = []
for i in range(descs_total):
desc_files.append(str(input()).rstrip())
return k, query_img, imgs_total, dist_f, desc_files
def load_concatenated_descriptors(desc_files, imgs_total):
#opening each descriptor file
descs_total = len(desc_files)
files = []
for i in range(descs_total):
files.append(open(desc_files[i]))
#loading dataset description that contains the name of each image
# and the feature vectors of each image and concatenating him
D = {}
print("Loading features from dataset")
for i in tq.tqdm(range(imgs_total)):
aux = ''
for j in range(1,descs_total):
aux += files[j].readline()[0:-1] + ','
#adding new feature vector to list
D[str(files[0].readline()[0:-1])] = list(map(int, aux[0:-1].split(',')))
#closing each descriptor file
for i in range(descs_total):
files[i].close()
return D
#recovering the concatenated descriptors of query image
def query_img_desc(QImg, DFiles):
H = []
print("Extracting features from query image")
for i in tq.tqdm(range(1, len(DFiles))):
desc = DFiles[i].split('/')[-1].split('.')[0]
img = im.imread(QImg)
if desc == 'BIC':
H += feb.BIC(img, 64)
elif desc == 'LBP':
H += feb.LPB(img)
return np.asarray(H).astype(np.float64)
def euclidean_dist(A, B):
return np.sum( (A - B)**2 ) ** 0.5
def distance(A, B, type):
if (type == 'euclidean'):
return euclidean_dist(A, B)
def toClass(ImgPath):
a = str(ImgPath.split('/')[-1])
b = a.split('_')
c = str(b[-1]).split('.')[0]
d = [b[1], c]
#d = [b[0], b[1], c]
tClass = '_'.join(d)
return tClass
#start############################################################################
# Main Function responsible for:
# 1) load input parameters
# 2) execute KNN
# 3) return the k similar images from query image
# 4) calc the precision and revocation metric
def main():
k, ImgQ_path, imgs_total, F, DFiles = load_parameters()
V = load_concatenated_descriptors(DFiles, imgs_total)
QImg = query_img_desc(ImgQ_path, DFiles)
d = {}
res = {}
H = {}
print("Calculating euclidean distance between query image and other images")
for ind, val in tq.tqdm(V.items()):
d[ind] = distance(QImg, val, F)
res = sorted(d.items(), key=lambda t: t[1])[1:k+1]
img = im.imread(ImgQ_path)
plt.figure(figsize=(2*k,2*k))
plt.subplots_adjust(wspace=0.5, hspace=0.5)
plt.subplot(k+1, 1, 1)
plt.axis('off')
plt.title('Query Image: '+ toClass(ImgQ_path))
plt.imshow(img, cmap='gray')
i = 1
imgs_path = '/'.join(ImgQ_path.split('/')[0:-1]) + '/'
for ind, val in res:
imgname = imgs_path+str(ind)+'.jpg'
classImg = str(toClass(imgname))
img = im.imread(imgname)
plt.subplot(k+1, 1, i+1)
plt.axis('off')
plt.title(str(i)+'ª Similar Image: '+ classImg)
plt.imshow(img, cmap='gray')
if (H.get(classImg) == None):
H[classImg] = 1
else:
H[classImg] += 1
i += 1
classPredict = sorted(H.items(), key=lambda t: t[1])[-1]
print("Class: " + str(classPredict[0]))
print("Precision: " + str(classPredict[1] / k))
plt.show()
#end##############################################################################
# Execution of Program
main()