cli suite working

1 files changed, 41 insertions, 0 deletions

diff --git a/check/commands/phash/test.py b/check/commands/phash/test.py
new file mode 100644
index 0000000..7fe2ae3
--- /dev/null
+++ b/check/commands/phash/test.py
@@ -0,0 +1,41 @@
+"""
+Query the database with a test set
+"""
+
+import click
+import os
+import glob
+
+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_input_glob',
+  required=True,
+  help="Input glob to search -- e.g. '../docs/images/*.jpg'")
+@click.pass_context
+def cli(ctx, opt_input_glob):
+  """
+  Query the database with a test set
+  """
+  for fn in sorted(glob.iglob(opt_input_glob)):
+    im = Image.open(fn).convert('RGB')
+    phash = compute_phash_int(im)
+
+    hash = sha256(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(fn, hash_result, phash_result))