From 77557082d48a16f491d0e8e77554afe81855b9b7 Mon Sep 17 00:00:00 2001
From: littletomatodonkey <dazhiningsibuqu@163.com>
Date: Fri, 2 Jul 2021 18:48:37 +0800
Subject: [PATCH] add support to append index (#1006)

---
 deploy/configs/build_cartoon.yaml |  1 +
 deploy/configs/build_logo.yaml    |  1 +
 deploy/configs/build_product.yaml |  1 +
 deploy/configs/build_vehicle.yaml |  1 +
 deploy/python/build_gallery.py    |  3 ++-
 deploy/vector_search/interface.py | 23 +++++++++++++++++++++--
 6 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/deploy/configs/build_cartoon.yaml b/deploy/configs/build_cartoon.yaml
index 3c93a6fa..f739cde3 100644
--- a/deploy/configs/build_cartoon.yaml
+++ b/deploy/configs/build_cartoon.yaml
@@ -31,6 +31,7 @@ IndexProcess:
   index_path: "./recognition_demo_data_v1.0/gallery_cartoon/index/"
   image_root: "./recognition_demo_data_v1.0/gallery_cartoon/"
   data_file:  "./recognition_demo_data_v1.0/gallery_cartoon/data_file.txt"
+  append_index: False
   delimiter: "\t"
   dist_type: "IP"
   pq_size: 100
diff --git a/deploy/configs/build_logo.yaml b/deploy/configs/build_logo.yaml
index 1f3800e2..a806ec00 100644
--- a/deploy/configs/build_logo.yaml
+++ b/deploy/configs/build_logo.yaml
@@ -29,6 +29,7 @@ IndexProcess:
   index_path: "./recognition_demo_data_v1.0/gallery_logo/index/"
   image_root: "./recognition_demo_data_v1.0/gallery_logo/"
   data_file:  "./recognition_demo_data_v1.0/gallery_logo/data_file.txt"
+  append_index: False
   delimiter: "\t"
   dist_type: "IP"
   pq_size: 100
diff --git a/deploy/configs/build_product.yaml b/deploy/configs/build_product.yaml
index 1ae4c0d5..679a1a67 100644
--- a/deploy/configs/build_product.yaml
+++ b/deploy/configs/build_product.yaml
@@ -29,6 +29,7 @@ IndexProcess:
   index_path: "./recognition_demo_data_v1.0/gallery_product/index"
   image_root: "./recognition_demo_data_v1.0/gallery_product/"
   data_file:  "./recognition_demo_data_v1.0/gallery_product/data_file.txt"
+  append_index: False
   delimiter: "\t"
   dist_type: "IP"
   pq_size: 100
diff --git a/deploy/configs/build_vehicle.yaml b/deploy/configs/build_vehicle.yaml
index 4897f24a..e149d938 100644
--- a/deploy/configs/build_vehicle.yaml
+++ b/deploy/configs/build_vehicle.yaml
@@ -29,6 +29,7 @@ IndexProcess:
   index_path: "./recognition_demo_data_v1.0/gallery_vehicle/index/"
   image_root: "./recognition_demo_data_v1.0/gallery_vehicle/"
   data_file:  "./recognition_demo_data_v1.0/gallery_vehicle/data_file.txt"
+  append_index: False
   delimiter: "\t"
   dist_type: "IP"
   pq_size: 100
diff --git a/deploy/python/build_gallery.py b/deploy/python/build_gallery.py
index 2087d9e0..142e3cf2 100644
--- a/deploy/python/build_gallery.py
+++ b/deploy/python/build_gallery.py
@@ -86,7 +86,8 @@ class GalleryBuilder(object):
             gallery_vectors=gallery_features,
             gallery_docs=gallery_docs,
             pq_size=config['pq_size'],
-            index_path=config['index_path'])
+            index_path=config['index_path'],
+            append_index=config["append_index"])
 
 
 def main(config):
diff --git a/deploy/vector_search/interface.py b/deploy/vector_search/interface.py
index fd217304..8dcd86f5 100644
--- a/deploy/vector_search/interface.py
+++ b/deploy/vector_search/interface.py
@@ -132,7 +132,8 @@ class Graph_Index(object):
               gallery_vectors,
               gallery_docs=[],
               pq_size=100,
-              index_path='graph_index/'):
+              index_path='graph_index/',
+              append_index=False):
         """
         build index 
         """
@@ -181,7 +182,25 @@ class Graph_Index(object):
         self.gallery_doc_dict["dist_type"] = self.dist_type
         self.gallery_doc_dict["with_attr"] = self.with_attr
 
-        with open(index_path + "/info.json", "w") as f:
+        output_path = os.path.join(index_path, "info.json")
+        if append_index is True and os.path.exists(output_path):
+            with open(output_path, "r") as fin:
+                lines = fin.readlines()[0]
+                ori_gallery_doc_dict = json.loads(lines)
+            assert ori_gallery_doc_dict["dist_type"] == self.gallery_doc_dict[
+                "dist_type"]
+            assert ori_gallery_doc_dict["dim"] == self.gallery_doc_dict["dim"]
+            assert ori_gallery_doc_dict["with_attr"] == self.gallery_doc_dict[
+                "with_attr"]
+            offset = ori_gallery_doc_dict["total_num"]
+            for i in range(0, self.gallery_doc_dict["total_num"]):
+                ori_gallery_doc_dict[str(i + offset)] = self.gallery_doc_dict[
+                    str(i)]
+
+            ori_gallery_doc_dict["total_num"] += self.gallery_doc_dict[
+                "total_num"]
+            self.gallery_doc_dict = ori_gallery_doc_dict
+        with open(output_path, "w") as f:
             json.dump(self.gallery_doc_dict, f)
 
         print("finished creating index ...")
-- 
GitLab