test API

1 files changed, 52 insertions, 7 deletions

diff --git a/check/app/server/api.py b/check/app/server/api.py
index c4f9f80..322d899 100644
--- a/check/app/server/api.py
+++ b/check/app/server/api.py
@@ -2,11 +2,13 @@ import os
 import re
 import time
 import numpy as np
+import logging
 from flask import Blueprint, request, jsonify
 from PIL import Image
 
 from app.models.sql_factory import search_by_phash, add_phash
-from app.utils.im_utils import pil2np
+from app.utils.im_utils import compute_phash_int
+from app.utils.file_utils import sha256_stream
 
 sanitize_re = re.compile('[\W]+')
 valid_exts = ['.gif', '.jpg', '.jpeg', '.png']
@@ -22,29 +24,72 @@ def index():
   """
   return jsonify({ 'status': 'ok' })
 
-@api.route('/v1/match/', methods=['POST'])
-def upload():
+@api.route('/v1/match', methods=['POST'])
+def match():
   """
   Search by uploading an image
   """
   start = time.time()
+  logging.debug(start)
 
-  file = request.files['query_img']
+  file = request.files['q']
   fn = file.filename
   if fn.endswith('blob'): # FIX PNG IMAGES?
     fn = 'filename.jpg'
+  logging.debug(fn)
 
   basename, ext = os.path.splitext(fn)
   if ext.lower() not in valid_exts:
     return jsonify({
+      'success': False,
+      'match': False,
       'error': 'not_an_image'
     })
 
+  ext = ext[1:].lower()
+
   im = Image.open(file.stream).convert('RGB')
   phash = compute_phash_int(im)
 
-  threshold = request.args.get('threshold') || 6
+  logging.debug(phash)
+  try:
+    threshold = int(request.args.get('threshold') or 6)
+    limit = int(request.args.get('limit') or 1)
+    add = str(request.args.get('add') or 'true') == 'true'
+  except:
+    return jsonify({
+      'success': False,
+      'match': False,
+      'error': 'param_error'
+    })
+
+  results = search_by_phash(phash=phash, threshold=threshold, limit=limit)
 
-  res = search_by_phash(phash, threshold)
+  if len(results) == 0:
+    if add:
+      hash = sha256_stream(file)
+      add_phash(sha256=hash, phash=phash, ext=ext)
+    if limit == 1:
+      return jsonify({
+        'success': True,
+        'match': False,
+      })
+    else:
+      return jsonify({
+        'success': True,
+        'match': False,
+        'results': [],
+      })
+
+  if limit > 1:
+    return jsonify({
+      'success': True,
+      'match': True,
+      'results': results,
+    })
 
-  return jsonify({ 'res': res })
+  return jsonify({
+    'success': True,
+    'match': True,
+    'closest_match': results[0],
+  })