update save function

face_similarity/diary_cover_similarity.py

@@ -18,14 +18,14 @@ facerec_model_path = os.path.join(model_dir, "dlib_face_recognition_resnet_model
 shape_model_path = os.path.join(model_dir, "shape_predictor_68_face_landmarks.dat")
 faiss_index_path = os.path.join(base_dir, "_index", "diary_cover.index")
 
-diary_after_cover_vec_file = "./diary_after_cover_vec.txt"
-
 face_rec = dlib.face_recognition_model_v1(facerec_model_path)
 face_detector = dlib.get_frontal_face_detector()
 shape_predictor = dlib.shape_predictor(shape_model_path)
 FACE_TO_VEC_FUN = lambda img: face_to_vec(img, face_rec, face_detector, shape_predictor)
 FAISS_DIARY_INDEX = faiss.read_index(faiss_index_path)
 
+DIARY_AFTER_COVER_FEATURE_KEY = "strategy_embedding:diary:cover:after"
+
 
 @bind("strategy_embedding/face_similarity/hello")
 def hello():
@@ -93,31 +93,26 @@ def save_diary_image_info():
             faces = FACE_TO_VEC_FUN(img)
             for face in faces:
                 after_res_dict[diary_id] = face["feature"]
-    redis_client_db.hmset("strategy_embedding:diary:cover:after", after_res_dict)
+    redis_client_db.hmset(DIARY_AFTER_COVER_FEATURE_KEY, after_res_dict)
 
 
-def save_faiss_index(load_file, save_path):
-    with open(load_file, "r") as f:
-        ids = []
-        features = []
-        lines = f.readlines()
-        print("lines: " + str(len(lines)))
-        count = 0
-        for line in lines:
-            count += 1
-            tmp = line.split("\t")
-            ids.append(tmp[0])
-            features.append(np.array(json.loads(tmp[1])))
-            print("{} {}".format(count, tmp[0]))
+def save_faiss_index(save_path):
+    data = redis_client_db.hgetall(DIARY_AFTER_COVER_FEATURE_KEY)
+    ids = []
+    features = []
+    for (k, v) in data.items():
+        ids.append(str(k, "utf-8"))
+        features.append(np.array(json.loads(v)))
 
-        ids_np = np.array(ids).astype("int")
-        features_np = np.array(features).astype("float32")
-        index = faiss.IndexHNSWFlat(128, 32)
-        print("trained: " + str(index.is_trained))
-        index2 = faiss.IndexIDMap(index)
-        index2.add_with_ids(features_np, ids_np)
-        faiss.write_index(index2, save_path)
-        print("faiss index saved")
+    print("ids: " + str(len(ids)))
+    ids_np = np.array(ids).astype("int")
+    features_np = np.array(features).astype("float32")
+    index = faiss.IndexHNSWFlat(128, 32)
+    print("trained: " + str(index.is_trained))
+    index2 = faiss.IndexIDMap(index)
+    index2.add_with_ids(features_np, ids_np)
+    faiss.write_index(index2, save_path)
+    print("faiss index saved")
 
 
 @bind("strategy_embedding/face_similarity/diary_url")
@@ -179,10 +174,11 @@ def get_similar_diary_ids_by_face_features(feature, limit=0.1):
 #     redis_key = "doris:diary:face_similary"
 #     redis_client3.hmset(redis_key, res_dict)
 
+if __name__ == "__main__":
+    begin_time = time.time()
 
-def main():
     save_diary_image_info()
-    # save_faiss_index(diary_after_cover_vec_file, faiss_index_path)
+    # save_faiss_index(faiss_index_path)
 
     # imgs = [
     #     "https://pic.igengmei.com/2020/07/03/1437/1b9975bb0b81-w", "https://pic.igengmei.com/2020/07/01/1812/ca64827a83da-w",
@@ -224,10 +220,4 @@ def main():
     # res = get_similar_diary_ids_by_face_features(a)
     # print(res)
 
-
-if __name__ == "__main__":
-    begin_time = time.time()
-
-    main()
-
     print("total cost: {:.2f}mins".format((time.time() - begin_time) / 60))