save result

src/diary_cover_similarity.py

@@ -120,9 +120,23 @@ def get_similar_diary_ids_by_url(url, index, face_to_vec_f, limit=0.1):
         return []
 
 
-def save_diary_similarity(file):
+def get_similar_diary_ids_by_face_features(feature, index, face_to_vec_f, limit=0.1):
+    feature = np.array(feature).astype("float32")
+    _sources, _ids = index.search(np.array([feature]), 10)
+    ids = _ids.flat
+    scores = _sources.flat
+    tmp = list(set(zip(ids, scores)))
+    res = []
+    for (id, score) in tmp:
+        if score >= limit:
+            res.append((id, score))
+    res.sort(key=lambda x: x[1], reverse=True)
+    return res
+
+
+def save_diary_similarity(load_file, index_path, face_to_vec_f):
     res_dict = {}
-    with open(file) as f:
+    with open(load_file) as f:
         lines = f.readlines()
         print("lines: " + str(len(lines)))
         count = 0
@@ -132,7 +146,7 @@ def save_diary_similarity(file):
             id = tmp[0]
             feature = np.array(json.loads(tmp[1]))
             print("{} {}".format(count, id))
-            tup_res = get_similar_diary_ids_by_face_features(feature, faiss_index, face_to_vec_f)
+            tup_res = get_similar_diary_ids_by_face_features(feature, index_path, face_to_vec_f)
             if tup_res:
                 res_dict[id] = tup_res
     print("done: " + str(len(res_dict)))
@@ -143,20 +157,6 @@ def save_diary_similarity(file):
     redis_client3.hmset(redis_key, res_dict)
 
 
-def get_similar_diary_ids_by_face_features(feature, index, face_to_vec_f, limit=0.1):
-    feature = np.array(feature).astype("float32")
-    _sources, _ids = index.search(np.array([feature]), 10)
-    ids = _ids.flat
-    scores = _sources.flat
-    tmp = list(set(zip(ids, scores)))
-    res = []
-    for (id, score) in tmp:
-        if score >= limit:
-            res.append((id, score))
-    res.sort(key=lambda x: x[1], reverse=True)
-    return res
-
-
 def main():
     base_dir = os.getcwd()
     print("base_dir: " + base_dir)
@@ -175,7 +175,11 @@ def main():
     # save_diary_image_info(diary_after_cover_vec_file, face_to_vec_f)
     # save_faiss_index(diary_after_cover_vec_file, faiss_index_path)
 
-    faiss_index = faiss.read_index(faiss_index_path)
+    faiss_index = faiss.read_index(
+        faiss_index_path,
+        faiss_index_path,
+        face_to_vec_f,
+    )
 
     # imgs = [
     #     "https://pic.igengmei.com/2020/07/03/1437/1b9975bb0b81-w", "https://pic.igengmei.com/2020/07/01/1812/ca64827a83da-w",
@@ -217,7 +221,7 @@ def main():
     # res = get_similar_diary_ids_by_face_features(a, faiss_index, face_to_vec_f)
     # print(res)
 
-    save_diary_similarity(diary_after_cover_vec_file)
+    save_diary_similarity(diary_after_cover_vec_file, faiss_index, face_to_vec_f)
 
 
 if __name__ == "__main__":