summaryrefslogtreecommitdiff
path: root/check
diff options
context:
space:
mode:
Diffstat (limited to 'check')
-rw-r--r--check/app/server/api.py20
-rw-r--r--check/commands/imagehash/query.py18
-rw-r--r--check/commands/imagehash/test.py2
3 files changed, 37 insertions, 3 deletions
diff --git a/check/app/server/api.py b/check/app/server/api.py
index 12d955b..6b97b59 100644
--- a/check/app/server/api.py
+++ b/check/app/server/api.py
@@ -7,6 +7,7 @@ from PIL import Image
# from app.utils.im_utils import pil2np
from app.models.sql_factory import search_by_phash, add_phash
+from app.utils.im_utils import pil2np
sanitize_re = re.compile('[\W]+')
valid_exts = ['.gif', '.jpg', '.jpeg', '.png']
@@ -24,5 +25,22 @@ def index():
def upload():
"""Search by image"""
start = time.time()
+
+ file = request.files['query_img']
+ fn = file.filename
+ if fn.endswith('blob'): # FIX PNG IMAGES?
+ fn = 'filename.jpg'
+
+ basename, ext = os.path.splitext(fn)
+ if ext.lower() not in valid_exts:
+ return jsonify({
+ 'error': 'not_an_image'
+ })
+
+ im = Image.open(file.stream).convert('RGB')
+ im_np = pil2np(im)
+
+ res = search_by_phash(im_np)
+
# get threshold
- return jsonify({ 'status': 'ok' })
+ return jsonify({ 'res': res })
diff --git a/check/commands/imagehash/query.py b/check/commands/imagehash/query.py
index 4c51324..f5d3a54 100644
--- a/check/commands/imagehash/query.py
+++ b/check/commands/imagehash/query.py
@@ -3,8 +3,11 @@ Search the database for an image
"""
import click
+import os
-from app.models.sql_factory import search_by_phash, add_phash
+from app.models.sql_factory import search_by_phash
+from app.utils.im_utils import compute_phash_int
+from app.utils.file_utils import sha256
@click.command()
@click.option('-i', '--input', 'opt_fn',
@@ -16,3 +19,16 @@ def cli(ctx, opt_fn):
Search the database for an image
"""
print('Searching for a file...')
+
+ if not os.path.exists(opt_fn):
+ print("File does not exist")
+ return
+
+ hash = sha256(opt_fn)
+ phash = compute_phash_int(opt_fn)
+
+ res = search_by_hash(hash)
+ print("search by hash:", res)
+
+ res =search_by_phash(phash)
+ print("search by phash:", res)
diff --git a/check/commands/imagehash/test.py b/check/commands/imagehash/test.py
index 7a38ae6..b3ddbe5 100644
--- a/check/commands/imagehash/test.py
+++ b/check/commands/imagehash/test.py
@@ -15,4 +15,4 @@ def cli(ctx, opt_fn):
"""
Search the database for an image
"""
- print('Searching for a file...')
+ print('Query the database with a test set')