-
Notifications
You must be signed in to change notification settings - Fork 5
/
generate_index.py
69 lines (52 loc) · 1.85 KB
/
generate_index.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
# -*- coding: utf-8 -*-
# Author: zhleternity
import os
os.environ['KERAS_BACKEND']='theano'
import h5py
import numpy as np
import argparse
from extract_cnn_vgg16_keras import VGGNet
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
ap = argparse.ArgumentParser()
ap.add_argument("-database", required = True,
help = "Path to database which contains images to be indexed")
ap.add_argument("-index", required = True,
help = "Name of index file")
args = vars(ap.parse_args())
'''
Returns a list of filenames for all jpg images in a directory.
'''
def get_imlist(path):
return [os.path.join(path,f) for f in os.listdir(path) if f.endswith('.png')]
'''
Extract features and index the images
'''
if __name__ == "__main__":
db = args["database"]
img_list = get_imlist(db)
print(img_list)
print("--------------------------------------------------")
print(" feature extraction starts")
print("--------------------------------------------------")
feats = []
names = []
model = VGGNet()
for i, img_path in enumerate(img_list):
norm_feat = model.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"]
print("--------------------------------------------------")
print(" writing feature extraction results ...")
print("--------------------------------------------------")
h5f = h5py.File(output, 'w')
h5f.create_dataset('dataset_1', data = feats)
h5f.create_dataset('dataset_2', data = names)
# h5f.create_dataset('dataset_2', data = np.string_(names))
h5f.close()