modify some format problem
@@ -28,6 +28,7 @@ from python.predict_rec import RecPredictor
from utils import logger
from utils import config
+
def split_datafile(data_file, image_root, delimiter="\t"):
'''
data_file: image path and info, which can be splitted by spacer
@@ -70,7 +71,6 @@ class GalleryBuilder(object):
# when remove data in index, do not need extract fatures
if operation_method != "remove":
gallery_features = self._extract_features(gallery_images, config)
-
assert operation_method in [
"new", "remove", "append"
], "Only append, remove and new operation are supported"
@@ -114,10 +114,11 @@ class GalleryBuilder(object):
#build index
if config["dist_type"] == "hamming":
- index = faiss.index_binary_factory(config["embedding_size"], index_method)
+ index = faiss.index_binary_factory(config["embedding_size"],
+ index_method)
else:
- index = faiss.index_factory(config["embedding_size"], index_method,
- dist_type)
+ index = faiss.index_factory(config["embedding_size"],
+ index_method, dist_type)
index = faiss.IndexIDMap2(index)
ids = {}
@@ -158,24 +159,25 @@ class GalleryBuilder(object):
# store faiss index file and id_map file
if config["dist_type"] == "hamming":
- faiss.write_index_binary(index,
- os.path.join(config["index_dir"], "vector.index"))
+ faiss.write_index_binary(
+ index, os.path.join(config["index_dir"], "vector.index"))
else:
- faiss.write_index(index,
- os.path.join(config["index_dir"], "vector.index"))
+ faiss.write_index(
+ index, os.path.join(config["index_dir"], "vector.index"))
with open(os.path.join(config["index_dir"], "id_map.pkl"), 'wb') as fd:
pickle.dump(ids, fd)
def _extract_features(self, gallery_images, config):
-
# extract gallery features
if config["dist_type"] == "hamming":
gallery_features = np.zeros(
- [len(gallery_images), config['embedding_size'] // 8], dtype=np.uint8)
+ [len(gallery_images), config['embedding_size'] // 8],
+ dtype=np.uint8)
else:
gallery_features = np.zeros(
- [len(gallery_images), config['embedding_size']], dtype=np.float32)
+ [len(gallery_images), config['embedding_size']],
+ dtype=np.float32)
#construct batch imgs and do inference
batch_size = config.get("batch_size", 32)
