cli suite working

1 files changed, 45 insertions, 0 deletions

diff --git a/check/commands/phash/query.py b/check/commands/phash/query.py
new file mode 100644
index 0000000..8fc8c61
--- /dev/null
+++ b/check/commands/phash/query.py
@@ -0,0 +1,45 @@
+"""
+Search the database for an image
+"""
+
+import click
+import os
+
+from PIL import Image
+
+from app.models.sql_factory import search_by_phash, search_by_hash
+from app.utils.im_utils import compute_phash_int
+from app.utils.file_utils import sha256
+
+@click.command()
+@click.option('-i', '--input', 'opt_fn',
+  required=True,
+  help="File to search")
+@click.pass_context
+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
+
+  im = Image.open(opt_fn).convert('RGB')
+  phash = compute_phash_int(im)
+
+  hash = sha256(opt_fn)
+
+  phash_match = search_by_phash(phash)
+  hash_match = search_by_hash(hash)
+
+  hash_result = 'NO'
+  if hash_match:
+    hash_result = 'YES'
+
+  phash_result = 'NO'
+  if len(phash_match):
+    phash_result = 'YES, score={}'.format(phash_match[0]['score'])
+
+  print("{} - hash={}, phash={}".format(opt_fn, hash_result, phash_result))