params

1 files changed, 47 insertions, 20 deletions

diff --git a/check/app/server/api.py b/check/app/server/api.py
index 63de5b6..66a0dd1 100644
--- a/check/app/server/api.py
+++ b/check/app/server/api.py
@@ -15,7 +15,11 @@ from app.utils.file_utils import sha256_stream
 sanitize_re = re.compile('[\W]+')
 valid_exts = ['.gif', '.jpg', '.jpeg', '.png']
 
-LIMIT = 9
+MATCH_THRESHOLD = 20
+MATCH_LIMIT = 10
+
+SIMILAR_THRESHOLD = 20
+SIMILAR_LIMIT = 10
 
 api = Blueprint('api', __name__)
 
@@ -26,17 +30,10 @@ def index():
   """
   return jsonify({ 'status': 'ok' })
 
-@api.route('/v1/match', methods=['POST'])
-def match():
-  """
-  Search by uploading an image
-  """
-  start = time.time()
-
+def get_params(default_threshold=MATCH_THRESHOLD, default_limit=MATCH_LIMIT):
   try:
-    threshold = int(request.form.get('threshold') or 6)
-    limit = int(request.form.get('limit') or 1)
-    add = str(request.form.get('add') or 'true') == 'true'
+    threshold = int(request.form.get('threshold') or default_threshold)
+    limit = int(request.form.get('limit') or default_limit)
   except:
     return jsonify({
       'success': False,
@@ -59,6 +56,7 @@ def match():
         'error': 'not_an_image'
       })
 
+    raw = None
     im = Image.open(file.stream).convert('RGB')
   else:
     url = request.form.get('url')
@@ -81,13 +79,22 @@ def match():
     raw = remote_response.read()
     im = Image.open(io.BytesIO(raw)).convert('RGB')
 
+@api.route('/v1/match', methods=['POST'])
+def match():
+  """
+  Search by uploading an image
+  """
+  start = time.time()
+
+  threshold, limit, raw, im = get_params()
+
   phash = compute_phash_int(im)
   ext = ext[1:].lower()
 
   results = search_by_phash(phash=phash, threshold=threshold, limit=limit)
 
   if len(results) == 0:
-    if add and url:
+    if url:
       # hash = sha256_stream(file)
       hash = sha256_stream(io.BytesIO(raw))
       add_phash(sha256=hash, phash=phash, ext=ext, url=url)
@@ -97,17 +104,37 @@ def match():
 
   logging.debug('query took {0:.2g} s.'.format(time.time() - start))
 
-  if limit > 1:
-    return jsonify({
-      'success': True,
-      'match': match,
-      'results': results,
-      'timing': time.time() - start,
-    })
+  return jsonify({
+    'success': True,
+    'match': match,
+    'results': results,
+    'timing': time.time() - start,
+  })
+
+@api.route('/v1/similar', methods=['POST'])
+def similar():
+  """
+  Search by uploading an image
+  """
+  start = time.time()
+
+  threshold, limit, raw, im = get_params(default_threshold=SIMILARITY_THRESHOLD, default_limit=SIMILARITY_LIMIT)
+
+  phash = compute_phash_int(im)
+  ext = ext[1:].lower()
+
+  results = search_by_phash(phash=phash, threshold=threshold, limit=limit)
+
+  if len(results) == 0:
+    match = False
+  else:
+    match = True
+
+  logging.debug('query took {0:.2g} s.'.format(time.time() - start))
 
   return jsonify({
     'success': True,
     'match': match,
-    'closest_match': results[0] if len(results) else None,
+    'results': results,
     'timing': time.time() - start,
   })