Commit b5de32b6 authored by 赵威's avatar 赵威

add limit

parent 20ac613f
...@@ -95,7 +95,7 @@ def save_faiss_index(load_file, save_path): ...@@ -95,7 +95,7 @@ def save_faiss_index(load_file, save_path):
faiss.write_index(index2, save_path) faiss.write_index(index2, save_path)
def get_similar_diary_ids_by_url(url, index, face_to_vec_f, limit=0): def get_similar_diary_ids_by_url(url, index, face_to_vec_f, limit=0.1):
img = url_to_ndarray(url) img = url_to_ndarray(url)
if img.any(): if img.any():
faces = face_to_vec_f(img) faces = face_to_vec_f(img)
...@@ -104,7 +104,11 @@ def get_similar_diary_ids_by_url(url, index, face_to_vec_f, limit=0): ...@@ -104,7 +104,11 @@ def get_similar_diary_ids_by_url(url, index, face_to_vec_f, limit=0):
_scores, _ids = index.search(np.array([face_feature]), 10) _scores, _ids = index.search(np.array([face_feature]), 10)
ids = _ids.flat ids = _ids.flat
scores = _scores.flat scores = _scores.flat
res = list(zip(ids, scores)) tmp = list(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) res.sort(key=lambda x: x[1], reverse=True)
return res return res
else: else:
...@@ -133,7 +137,7 @@ def main(): ...@@ -133,7 +137,7 @@ def main():
img_url = "https://pic.igengmei.com/2020/07/03/1437/1b9975bb0b81-w" img_url = "https://pic.igengmei.com/2020/07/03/1437/1b9975bb0b81-w"
res = get_similar_diary_ids_by_url(img_url, faiss_index, face_to_vec_f) res = get_similar_diary_ids_by_url(img_url, faiss_index, face_to_vec_f, limit=0.18232107)
print(res) print(res)
# a = [ # a = [
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment