def trainh5(databasepath="",h5name=""):
database = databasepath
index = 'models/'+h5name+'.h5'
img_list = get_imlist(database)
print("--------------------------------------------------")
print(" 训练集开始 ")
print("--------------------------------------------------")
feats = []
names = []
model = VGGNet()
for i, img_path in enumerate(img_list):
norm_feat = model.vgg_extract_feat(img_path) #提取特征
img_name = os.path.split(img_path)[1]
feats.append(norm_feat)
names.append(img_name)
print("提取图片 No. %d 的特征,总共有 %d 张图片"
% ((i + 1), len(img_list)))
feats = np.array(feats)
output = index
print("--------------------------------------------------")
print(" 正在写入特征库 ... ")
print("--------------------------------------------------")
h5f = h5py.File(output, 'w')
h5f.create_dataset('dataset_1', data=feats)
h5f.create_dataset('dataset_2', data=np.string_(names))
h5f.close()
def search(querypath, h5name):
# 要搜的原图path，h5文件的name，原数据集的path
query = querypath
index = 'models/' + h5name + '.h5'
result = 'database/' + h5name + '/'
h5f = h5py.File(index, 'r')
feats = h5f['dataset_1'][:]
imgNames = h5f['dataset_2'][:]
h5f.close()
# read query image
queryImg = mpimg.imread(query)
# 初始化 VGGNet16 模块
model = VGGNet()
#提取query图片的特征
queryVec = model.vgg_extract_feat(query)
scores = np.dot(queryVec, feats.T)
rank_ID = np.argsort(scores)[::-1]
rank_score = scores[rank_ID]
# 检索出三张相似度最高的图片
maxres = 11
imlist = []
# 输出相似度
for i, index in enumerate(rank_ID[0:maxres]):
imlist.append(imgNames[index])
#print("image names: " +
# str(imgNames[index]) +
# " scores: %f"
# % rank_score[i])
#print("top %d images in order are: " % maxres, imlist)
# 输出前三张
return imlist
def query(filename, k):
clear_session()
model = VGGNet()
queryVec = model.vgg_extract_feat(filename)
p = hnswlib.Index('l2', 512)
p.load_index('hnswindex.bin')
p.set_ef(50) #p>k
return p.knn_query(queryVec, k)
'''
if __name__ == "__main__":
database = '../data/picture'
index = 'vgg_featureCNN.h5'
img_list = get_imlist(database)
print("--------------------------------------------------")
print(" feature extraction starts")
print("--------------------------------------------------")
feats = []
names = []
model = VGGNet()
for i, img_path in enumerate(img_list):
norm_feat = model.vgg_extract_feat(img_path) #修改此处改变提取特征的网络
img_name = os.path.split(img_path)[1]
feats.append(norm_feat)
names.append(img_name)
print("extracting feature from image No. %d , %d images in total" %((i+1), len(img_list)))
feats = np.array(feats)
# print(feats)
# directory for storing extracted features
# output = args["index"]
output = index
print("--------------------------------------------------")
print(" writing feature extraction results ...")
print("--------------------------------------------------")
print(feats)
imgNames = h5f['dataset_2'][...]
h5f.close()
print("--------------------------------------------------")
print(" searching starts")
print("--------------------------------------------------")
# read and show query image
# queryDir = args["query"]
# init VGGNet16 model
model = VGGNet()
# extract query image's feature, compute simlarity score and sort
queryVec = model.vgg_extract_feat(query) #修改此處改變提取特徵的網路
print(queryVec.shape)
print(feats.shape)
scores = np.dot(queryVec, feats.T)
rank_ID = np.argsort(scores)[::-1]
rank_score = scores[rank_ID]
# print (rank_ID)
print(rank_score)
# number of top retrieved images to show
maxres = 3 #檢索出三張相似度最高的圖片
imlist = []
for i, index in enumerate(rank_ID[0:maxres]):
imlist.append(imgNames[index])
# print(type(imgNames[index]))
print("image names: " + str(imgNames[index]) +
print("--------------------------------------------------")
print(" searching starts")
print("--------------------------------------------------")
# read and show query image
queryImg = mpimg.imread(query)
plt.title("Query Image")
plt.imshow(queryImg)
plt.show()
# init VGGNet16 model
model = VGGNet()
# extract query image's feature, compute simlarity score and sort
queryVec = model.vgg_extract_feat(query) # 修改此处改变提取特征的网络
# print(queryVec.shape)
# print(feats.shape)
print('--------------------------')
# print(queryVec)
# print(feats.T)
print('--------------------------')
scores = np.dot(queryVec, feats.T)
# scores = np.dot(queryVec, feats.T)/(np.linalg.norm(queryVec)*np.linalg.norm(feats.T))
rank_ID = np.argsort(scores)[::-1]
rank_score = scores[rank_ID]
# print (rank_ID)
print(rank_score)
# number of top retrieved images to show
maxres = 5 # 检索出五张相似度最高的图片
'''
if __name__ == "__main__":
database = 'pic/'
index = 'vgg_featureCNN.h5'
img_list = get_imlist(database)
print("--------------------------------------------------")
print(" feature extraction starts")
print("--------------------------------------------------")
feats = []
names = []
model = VGGNet()
for i, img_path in enumerate(img_list):
norm_feat = model.vgg_extract_feat(img_path) #修改此處改變提取特徵的網路
img_name = os.path.split(img_path)[1]
feats.append(norm_feat)
names.append(img_name)
print("extracting feature from image No. %d , %d images in total" %((i+1), len(img_list)))
feats = np.array(feats)
# print(feats)
# directory for storing extracted features
# output = args["index"]
output = index
print("--------------------------------------------------")
print(" writing feature extraction results ...")
print("--------------------------------------------------")