@@ -29,11 +29,12 @@ class IndexData:
29
29
IDMAP_FN = "id_map"
30
30
IDMAP_SUFFIX = ".yaml"
31
31
32
- def __init__ (self , index , id_map , metric_type , index_type ):
32
+ def __init__ (self , index , index_info ):
33
33
self ._index = index
34
- self ._id_map = id_map
35
- self ._metric_type = metric_type
36
- self ._index_type = index_type
34
+ self ._index_info = index_info
35
+ self ._id_map = index_info ["id_map" ]
36
+ self ._metric_type = index_info ["metric_type" ]
37
+ self ._index_type = index_info ["index_type" ]
37
38
38
39
@property
39
40
def index (self ):
@@ -260,7 +261,9 @@ def build(
260
261
index , ids = cls ._add_gallery (
261
262
metric_type , index , ids , features , gallery_docs , mode = "new"
262
263
)
263
- return IndexData (index , ids , metric_type , index_type )
264
+ return IndexData (
265
+ index , {"id_map" : ids , "metric_type" : metric_type , "index_type" : index_type }
266
+ )
264
267
265
268
@classmethod
266
269
def remove (
@@ -288,7 +291,9 @@ def remove(
288
291
# remove ids in id_map, remove index data in faiss index
289
292
index .remove_ids (remove_ids )
290
293
ids = {k : v for k , v in ids .items () if k not in remove_ids }
291
- return IndexData (index , ids , metric_type , index_type )
294
+ return IndexData (
295
+ index , {"id_map" : ids , "metric_type" : metric_type , "index_type" : index_type }
296
+ )
292
297
293
298
@classmethod
294
299
def append (cls , gallery_imgs , gallery_label , predict_func , index ):
@@ -310,7 +315,9 @@ def append(cls, gallery_imgs, gallery_label, predict_func, index):
310
315
index , ids = cls ._add_gallery (
311
316
metric_type , index , ids , features , gallery_docs , mode = "append"
312
317
)
313
- return IndexData (index , ids , metric_type , index_type )
318
+ return IndexData (
319
+ index , {"id_map" : ids , "metric_type" : metric_type , "index_type" : index_type }
320
+ )
314
321
315
322
@classmethod
316
323
def _add_gallery (
0 commit comments