save res

src/diary_cover_similarity.py

@@ -7,9 +7,9 @@ import dlib
 import faiss
 import numpy as np
 
+from utils.cache import redis_client3
 from utils.es import es_query
 from utils.images import face_to_vec, url_to_ndarray
-from utils.cache import redis_client3
 
 
 def save_diary_image_info(save_file, face_to_vec_f):
@@ -120,6 +120,29 @@ def get_similar_diary_ids_by_url(url, index, face_to_vec_f, limit=0.1):
         return []
 
 
+def save_diary_similarity(file):
+    res_dict = {}
+    with open(file) as f:
+        lines = f.readlines()
+        print("lines: " + str(len(lines)))
+        count = 0
+        for line in lines:
+            count += 1
+            tmp = line.split("\t")
+            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)
+            if tup_res:
+                res_dict[id] = tup_res
+    print("done: " + str(len(res_dict)))
+    key = random.choice(list(res_dict.keys()))
+    print(res_dict[key])
+
+    redis_key = "doris:diary:face_similary"
+    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)
@@ -194,26 +217,7 @@ def main():
     # res = get_similar_diary_ids_by_face_features(a, faiss_index, face_to_vec_f)
     # print(res)
 
-    res_dict = {}
-    with open(diary_after_cover_vec_file) as f:
-        lines = f.readlines()
-        print("lines: " + str(len(lines)))
-        count = 0
-        for line in lines:
-            count += 1
-            tmp = line.split("\t")
-            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)
-            if tup_res:
-                res_dict[id] = tup_res
-    print("done: " + str(len(res_dict)))
-    key = random.choice(list(res_dict.keys()))
-    print(res_dict[key])
-
-    redis_key = "doris:diary:face_similary"
-    redis_client3.hmset(redis_key, res_dict)
+    save_diary_similarity(diary_after_cover_vec_file)
 
 
 if __name__ == "__main__":