client,server: reverse search post by signature

2 files changed, 39 insertions, 6 deletions

diff --git a/server/szurubooru/api/post_api.py b/server/szurubooru/api/post_api.py
index 2134b3b..3a35bf8 100644
--- a/server/szurubooru/api/post_api.py
+++ b/server/szurubooru/api/post_api.py
@@ -14,7 +14,7 @@ from szurubooru.func import (
     similar,
     snapshots,
     tags,
-    versions,
+    versions, image_hash,
 )
 
 _search_executor_config = search.configs.PostSearchConfig()
@@ -321,6 +321,31 @@ def get_posts_by_image(
     }
 
 
+@rest.routes.get("/post/(?P<post_id>[^/]+)/reverse-search/?")
+def get_posts_lookalikes(
+        ctx: rest.Context, params: Dict[str, str] = {}
+) -> rest.Response:
+    auth.verify_privilege(ctx.user, "posts:reverse_search")
+    limit = ctx.get_param_as_int("limit", default=10, min=1, max=100)
+    threshold = ctx.get_param_as_float("threshold", default=1, min=0, max=100)
+    post_id = _get_post_id(params)
+    post = posts.get_post_by_id(post_id)
+    sig = image_hash.unpack_signature(post.signature.signature)
+    lookalikes = posts.search_by_signature(sig, limit, threshold)
+    # exclude the same post:
+    lookalikes = filter(lambda la: la[1].post_id != post_id, lookalikes)
+    lookalikes = sorted(lookalikes, key=lambda la: la[0])
+    return {
+        "similarPosts": [
+            {
+                "distance": distance,
+                "post": _serialize_post(ctx, post),
+            }
+            for distance, post in lookalikes
+        ],
+    }
+
+
 @rest.routes.get("/posts/median/?")
 def get_posts_median(
     ctx: rest.Context, _params: Dict[str, str] = {}
diff --git a/server/szurubooru/func/posts.py b/server/szurubooru/func/posts.py
index 107fc5d..fac7cd5 100644
--- a/server/szurubooru/func/posts.py
+++ b/server/szurubooru/func/posts.py
@@ -21,6 +21,7 @@ from szurubooru.func import (
     users,
     util,
 )
+from szurubooru.func.image_hash import NpMatrix
 
 logger = logging.getLogger(__name__)
 
@@ -954,8 +955,15 @@ def search_by_image_exact(image_content: bytes) -> Optional[model.Post]:
 
 def search_by_image(image_content: bytes) -> List[Tuple[float, model.Post]]:
     query_signature = image_hash.generate_signature(image_content)
-    query_words = image_hash.generate_words(query_signature)
+    return search_by_signature(query_signature)
 
+
+def search_by_signature(
+    signature: NpMatrix,
+    limit: int = 100,
+    distance_cutoff: float = image_hash.DISTANCE_CUTOFF,
+) -> List[Tuple[float, model.Post]]:
+    query_words = image_hash.generate_words(signature)
     """
     The unnest function is used here to expand one row containing the 'words'
     array into multiple rows each containing a singular word.
@@ -969,10 +977,10 @@ def search_by_image(image_content: bytes) -> List[Tuple[float, model.Post]]:
     FROM post_signature AS s, unnest(s.words, :q) AS a(word, query)
     WHERE a.word = a.query
     GROUP BY s.post_id
-    ORDER BY score DESC LIMIT 100;
+    ORDER BY score DESC LIMIT :limit;
     """
 
-    candidates = db.session.execute(dbquery, {"q": query_words})
+    candidates = db.session.execute(dbquery, {"q": query_words, "limit": limit})
     data = tuple(
         zip(
             *[
@@ -983,13 +991,13 @@ def search_by_image(image_content: bytes) -> List[Tuple[float, model.Post]]:
     )
     if data:
         candidate_post_ids, sigarray = data
-        distances = image_hash.normalized_distance(sigarray, query_signature)
+        distances = image_hash.normalized_distance(sigarray, signature)
         return [
             (distance, try_get_post_by_id(candidate_post_id))
             for candidate_post_id, distance in zip(
                 candidate_post_ids, distances
             )
-            if distance < image_hash.DISTANCE_CUTOFF
+            if distance < distance_cutoff
         ]
     else:
         return []