2 files changed, 108 insertions, 18 deletions

diff --git a/check/app/server/api.py b/check/app/server/api.py
index c4f9f80..66a0dd1 100644
--- a/check/app/server/api.py
+++ b/check/app/server/api.py
@@ -1,17 +1,25 @@
+import io
 import os
 import re
 import time
 import numpy as np
+import logging
+import urllib.request
 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']
 
-LIMIT = 9
+MATCH_THRESHOLD = 20
+MATCH_LIMIT = 10
+
+SIMILAR_THRESHOLD = 20
+SIMILAR_LIMIT = 10
 
 api = Blueprint('api', __name__)
 
@@ -22,29 +30,111 @@ def index():
   """
   return jsonify({ 'status': 'ok' })
 
-@api.route('/v1/match/', methods=['POST'])
-def upload():
+def get_params(default_threshold=MATCH_THRESHOLD, default_limit=MATCH_LIMIT):
+  try:
+    threshold = int(request.form.get('threshold') or default_threshold)
+    limit = int(request.form.get('limit') or default_limit)
+  except:
+    return jsonify({
+      'success': False,
+      'match': False,
+      'error': 'param_error'
+    })
+
+  if 'q' in request.files:
+    file = request.files['q']
+    fn = file.filename
+    if fn.endswith('blob'): # FIX PNG IMAGES?
+      logging.debug('received a blob, assuming JPEG')
+      fn = 'filename.jpg'
+
+    basename, ext = os.path.splitext(fn)
+    if ext.lower() not in valid_exts:
+      return jsonify({
+        'success': False,
+        'match': False,
+        'error': 'not_an_image'
+      })
+
+    raw = None
+    im = Image.open(file.stream).convert('RGB')
+  else:
+    url = request.form.get('url')
+    if not url:
+      return jsonify({
+        'success': False,
+        'match': False,
+        'error': 'no_image'
+      })
+    basename, ext = os.path.splitext(url)
+    if ext.lower() not in valid_exts:
+      return jsonify({
+        'success': False,
+        'match': False,
+        'error': 'not_an_image'
+      })
+
+    remote_request = urllib.request.Request(url)
+    remote_response = urllib.request.urlopen(remote_request)
+    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()
 
-  file = request.files['query_img']
-  fn = file.filename
-  if fn.endswith('blob'): # FIX PNG IMAGES?
-    fn = 'filename.jpg'
+  threshold, limit, raw, im = get_params()
 
-  basename, ext = os.path.splitext(fn)
-  if ext.lower() not in valid_exts:
-    return jsonify({
-      'error': 'not_an_image'
-    })
+  phash = compute_phash_int(im)
+  ext = ext[1:].lower()
+
+  results = search_by_phash(phash=phash, threshold=threshold, limit=limit)
+
+  if len(results) == 0:
+    if url:
+      # hash = sha256_stream(file)
+      hash = sha256_stream(io.BytesIO(raw))
+      add_phash(sha256=hash, phash=phash, ext=ext, url=url)
+    match = False
+  else:
+    match = True
+
+  logging.debug('query took {0:.2g} s.'.format(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)
 
-  im = Image.open(file.stream).convert('RGB')
   phash = compute_phash_int(im)
+  ext = ext[1:].lower()
+
+  results = search_by_phash(phash=phash, threshold=threshold, limit=limit)
 
-  threshold = request.args.get('threshold') || 6
+  if len(results) == 0:
+    match = False
+  else:
+    match = True
 
-  res = search_by_phash(phash, threshold)
+  logging.debug('query took {0:.2g} s.'.format(time.time() - start))
 
-  return jsonify({ 'res': res })
+  return jsonify({
+    'success': True,
+    'match': match,
+    'results': results,
+    'timing': time.time() - start,
+  })
diff --git a/check/app/server/create.py b/check/app/server/create.py
index 788bfaa..6dab331 100644
--- a/check/app/server/create.py
+++ b/check/app/server/create.py
@@ -27,7 +27,7 @@ def create_app(script_info=None):
   """
   functional pattern for creating the flask app
   """
-  app = Flask(__name__, static_folder='static', static_url_path='')
+  app = Flask(__name__, static_folder='static', static_url_path='/static')
   app.config['SQLALCHEMY_DATABASE_URI'] = connection_url
   app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
   app.config['CELERY_BROKER_URL'] = cfg.CELERY_BROKER_URL