save all diaries

src/diary_cover_similarity.py

@@ -56,27 +56,29 @@ def save_diary_image_info(save_file, face_to_vec_f):
 
     with open(save_file, "w") as f:
         step = 100
-        # for i in range(0, 500000, step):
-        res_dict = es_query("diary", q, 0, step)
-
-        for item in res_dict["hits"]["hits"]:
-            diary_id = item["_source"]["id"]
-            before_cover_url = item["_source"]["before_cover_url"] + "-w"
-            after_cover_url = item["_source"]["after_cover_url"] + "-w"
-            img = url_to_ndarray(after_cover_url)
-            if img.any():
-                faces = face_to_vec_f(img)
-                for face in faces:
-                    line = str(diary_id) + "\t" + face["feature"] + "\n"
-                    # print(line)
-                    f.write(line)
+        for i in range(0, 500000, step):
+            res_dict = es_query("diary", q, 0, step)
+
+            for item in res_dict["hits"]["hits"]:
+                diary_id = item["_source"]["id"]
+                # before_cover_url = item["_source"]["before_cover_url"] + "-w"
+                after_cover_url = item["_source"]["after_cover_url"] + "-w"
+                img = url_to_ndarray(after_cover_url)
+                if img.any():
+                    faces = face_to_vec_f(img)
+                    for face in faces:
+                        line = str(diary_id) + "\t" + face["feature"] + "\n"
+                        # print(line)
+                        f.write(line)
 
 
 def save_faiss_index(load_file, save_path):
     with open(load_file, "r") as f:
         ids = []
         features = []
-        for line in f.readlines():
+        lines = f.readlines()
+        print("lines: " + str(len(lines)))
+        for line in lines:
             tmp = line.split("\t")
             ids.append(tmp[0])
             features.append(np.array(json.loads(tmp[1])))
@@ -105,8 +107,8 @@ def main():
     shape_predictor = dlib.shape_predictor(shape_model_path)
     face_to_vec_f = lambda img: face_to_vec(img, face_rec, face_detector, shape_predictor)
 
-    # save_diary_image_info(diary_after_cover_vec_file, face_to_vec_f)
-    # save_faiss_index(diary_after_cover_vec_file, faiss_index_path)
+    save_diary_image_info(diary_after_cover_vec_file, face_to_vec_f)
+    save_faiss_index(diary_after_cover_vec_file, faiss_index_path)
 
     a = [
         -0.08361373096704483, 0.06760436296463013, 0.10752949863672256, -0.020746365189552307, -0.07035162299871445,