save diary info

src/diary_cover_similarity.py

@@ -53,9 +53,10 @@ def save_diary_image_info(save_file, face_to_vec_f):
             "include": ["id", "before_cover_url", "after_cover_url"]
         }
     }
-
+    count = 0
     with open(save_file, "w") as f:
         step = 100
+
         for i in range(0, 500000, step):
             res_dict = es_query("diary", q, 0, step)
 
@@ -65,9 +66,11 @@ def save_diary_image_info(save_file, face_to_vec_f):
                 after_cover_url = item["_source"]["after_cover_url"] + "-w"
                 img = url_to_ndarray(after_cover_url)
                 if img.any():
+                    count += 1
+                    print("count: " + str(count))
                     faces = face_to_vec_f(img)
                     for face in faces:
-                        line = str(diary_id) + "\t" + face["feature"] + "\t" + after_cover_url + "\n"
+                        line = str(diary_id) + "\t" + face["feature"] + "\n"
                         # print(line)
                         f.write(line)
 
@@ -115,6 +118,20 @@ def get_similar_diary_ids_by_url(url, index, face_to_vec_f, limit=0.1):
         return []
 
 
+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)
+    print(res)
+
+
 def main():
     base_dir = os.getcwd()
     print("base_dir: " + base_dir)
@@ -131,17 +148,19 @@ def main():
     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_faiss_index(diary_after_cover_vec_file, faiss_index_path)
 
     faiss_index = faiss.read_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",
-        "https://pic.igengmei.com/2020/07/04/1711/24f4131a9b1e-w", "https://pic.igengmei.com/2020/07/04/1507/e17a995be219-w"
-    ]
-    for img_url in imgs:
-        res = get_similar_diary_ids_by_url(img_url, faiss_index, face_to_vec_f, limit=0.18232107)
-        print(res)
 
+    # imgs = [
+    #     "https://pic.igengmei.com/2020/07/03/1437/1b9975bb0b81-w", "https://pic.igengmei.com/2020/07/01/1812/ca64827a83da-w",
+    #     "https://pic.igengmei.com/2020/07/04/1711/24f4131a9b1e-w", "https://pic.igengmei.com/2020/07/04/1507/e17a995be219-w"
+    # ]
+    # for img_url in imgs:
+    #     res = get_similar_diary_ids_by_url(img_url, faiss_index, face_to_vec_f, limit=0.18232107)
+    #     print(res)
+
+    # print("@@@@@@@@")
     # a = [
     #     -0.08361373096704483, 0.06760436296463013, 0.10752949863672256, -0.020746365189552307, -0.07035162299871445,
     #     -0.014547230675816536, -0.043201886117458344, -0.12196271121501923, 0.13929598033428192, -0.1360183209180832,
@@ -170,13 +189,7 @@ def main():
     #     -0.015351934358477592, -0.04335442930459976, -0.26258283853530884, -0.021509556099772453, 0.12185295671224594,
     #     -0.011788002215325832, 0.01337978895753622, -0.008025042712688446
     # ]
-    # feature = np.array(a).astype("float32")
-    # index = faiss.read_index(faiss_index_path)
-    # D, I = index.search(np.array([feature]), 5)
-    # ids = I.flat
-    # scores = D.flat
-    # res = list(zip(ids, scores))
-    # res.sort(key = lambda x: x[1], reverse=True)
+    # res = get_similar_diary_ids_by_face_features(a)
     # print(res)