test API

1 files changed, 19 insertions, 23 deletions

diff --git a/check/commands/phash/query.py b/check/commands/phash/query.py
index 8fc8c61..7fe2ae3 100644
--- a/check/commands/phash/query.py
+++ b/check/commands/phash/query.py
@@ -1,9 +1,10 @@
 """
-Search the database for an image
+Query the database with a test set
 """
 
 import click
 import os
+import glob
 
 from PIL import Image
 
@@ -12,34 +13,29 @@ from app.utils.im_utils import compute_phash_int
 from app.utils.file_utils import sha256
 
 @click.command()
-@click.option('-i', '--input', 'opt_fn',
+@click.option('-i', '--input', 'opt_input_glob',
   required=True,
-  help="File to search")
+  help="Input glob to search -- e.g. '../docs/images/*.jpg'")
 @click.pass_context
-def cli(ctx, opt_fn):
+def cli(ctx, opt_input_glob):
   """
-  Search the database for an image
+  Query the database with a test set
   """
-  print('Searching for a file...')
+  for fn in sorted(glob.iglob(opt_input_glob)):
+    im = Image.open(fn).convert('RGB')
+    phash = compute_phash_int(im)
 
-  if not os.path.exists(opt_fn):
-    print("File does not exist")
-    return
+    hash = sha256(fn)
 
-  im = Image.open(opt_fn).convert('RGB')
-  phash = compute_phash_int(im)
+    phash_match = search_by_phash(phash)
+    hash_match = search_by_hash(hash)
 
-  hash = sha256(opt_fn)
+    hash_result = 'NO'
+    if hash_match:
+      hash_result = 'YES'
 
-  phash_match = search_by_phash(phash)
-  hash_match = search_by_hash(hash)
+    phash_result = 'NO'
+    if len(phash_match):
+      phash_result = 'YES, score={}'.format(phash_match[0]['score'])
 
-  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))
+    print("{} - hash={}, phash={}".format(fn, hash_result, phash_result))