| import os
|
| import numpy as np
|
| import faiss
|
|
|
| EXP = "josias_v1"
|
| ROOT = os.path.dirname(os.path.abspath(__file__))
|
| LOGDIR = os.path.join(ROOT, "logs", EXP)
|
|
|
| fea_path = os.path.join(LOGDIR, "total_fea.npy")
|
| assert os.path.exists(fea_path), f"Missing: {fea_path}"
|
|
|
| x = np.load(fea_path)
|
|
|
| x = np.asarray(x, dtype=np.float32)
|
| if x.ndim == 1:
|
| x = x.reshape(-1, 1)
|
|
|
| n, d = x.shape
|
| print("Loaded features:", (n, d), x.dtype)
|
|
|
| if n == 0:
|
| raise RuntimeError("total_fea.npy is empty (0 rows). Cannot build index.")
|
|
|
|
|
| nlist = 1058
|
| nprobe = 1
|
|
|
|
|
| quantizer = faiss.IndexFlatL2(d)
|
| index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
|
|
|
| print("Training IVF index...")
|
| index.train(x)
|
|
|
| print("Adding vectors...")
|
| index.add(x)
|
|
|
| index.nprobe = nprobe
|
|
|
| out_index = os.path.join(LOGDIR, f"added_IVF{nlist}_Flat_nprobe_{nprobe}_{EXP}_v2.index")
|
| faiss.write_index(index, out_index)
|
|
|
| print("Wrote index:", out_index)
|
| print("ntotal =", index.ntotal) |