File size: 767 Bytes
decbc4d |
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 |
import glob
import os
import json
import tqdm
import numpy as np
from multiprocessing import Pool
from autofaiss import build_index
def load_file(file):
return np.load(file)
in_dir = '/data0/ImageBindFeatures/cc3m/'
save_path = "cc3m_imagebind.npy"
files = glob.glob(in_dir + '*/*.npy')
pool = Pool(256)
results = pool.map(load_file, files)
results = np.stack(results, axis=0)
np.save(save_path, results)
with open('cc3m_imagebind_files.json', 'w') as f:
json.dump([x.replace(in_dir, '') for x in files], f)
# build index
build_index(embeddings="./",
index_path="knn.index",
index_infos_path="infos.json",
max_index_memory_usage = "32G",
current_memory_available = "100G",
metric_type='ip') |