cli suite working

9 files changed, 152 insertions, 108 deletions

diff --git a/check/commands/imagehash/add.py b/check/commands/imagehash/add.py
deleted file mode 100644
index 73f8b69..0000000
--- a/check/commands/imagehash/add.py
+++ /dev/null
@@ -1,34 +0,0 @@
-"""
-Add a file to the database
-"""
-
-import click
-import os
-
-from app.models.sql_factory import search_by_phash, add_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',
-  required=True,
-  help="File to add (gif/jpg/png)")
-@click.option('-u', '--upload', 'opt_upload', is_flag=True,
-  help='Whether to upload this file to S3')
-@click.pass_context
-def cli(ctx, opt_fn, opt_upload):
-  """
-  Add a single file
-  """
-  print('Adding 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)
-
-  dir, fn = os.path.split(opt_fn)
-  root, ext = os.path.splitext(fn)
-
-  add_phash(sha256=hash, phash=phash, ext=ext)
diff --git a/check/commands/imagehash/load.py b/check/commands/imagehash/load.py
deleted file mode 100644
index e61b751..0000000
--- a/check/commands/imagehash/load.py
+++ /dev/null
@@ -1,22 +0,0 @@
-"""
-Loop over a directory of images
- - Compute their phashes
- - Optionally upload them to s3?
-"""
-
-import click
-
-from app.models.sql_factory import search_by_phash, add_phash
-
-@click.command()
-@click.option('-i', '--input', 'opt_dir_fn',
-  required=True,
-  help="File to add (gif/jpg/png)")
-@click.option('-u', '--upload', 'opt_upload', is_flag=True,
-  help='Whether to upload this file to S3')
-@click.pass_context
-def cli(ctx, opt_dir_fn, opt_store):
-  """
-  Add a directory of images
-  """
-  print('Adding a directory...')
diff --git a/check/commands/imagehash/query.py b/check/commands/imagehash/query.py
deleted file mode 100644
index f5d3a54..0000000
--- a/check/commands/imagehash/query.py
+++ /dev/null
@@ -1,34 +0,0 @@
-"""
-Search the database for an image
-"""
-
-import click
-import os
-
-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',
-  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
-
-  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
deleted file mode 100644
index b3ddbe5..0000000
--- a/check/commands/imagehash/test.py
+++ /dev/null
@@ -1,18 +0,0 @@
-"""
-Query the database with a test set
-"""
-
-import click
-
-from app.models.sql_factory import search_by_phash, add_phash
-
-@click.command()
-@click.option('-i', '--input', 'opt_fn',
-  required=True,
-  help="Directory to search")
-@click.pass_context
-def cli(ctx, opt_fn):
-  """
-  Search the database for an image
-  """
-  print('Query the database with a test set')
diff --git a/check/commands/phash/add.py b/check/commands/phash/add.py
new file mode 100644
index 0000000..1565202
--- /dev/null
+++ b/check/commands/phash/add.py
@@ -0,0 +1,21 @@
+"""
+Add a file to the database
+"""
+
+import click
+
+from app.models.sql_factory import add_phash_by_filename
+
+@click.command()
+@click.option('-i', '--input', 'opt_fn',
+  required=True,
+  help="File to add (gif/jpg/png)")
+@click.option('-u', '--upload', 'opt_upload', is_flag=True,
+  help='Whether to upload this file to S3')
+@click.pass_context
+def cli(ctx, opt_fn, opt_upload):
+  """
+  Add a single file
+  """
+  print('Adding a file...')
+  add_phash_by_filename(opt_fn)
diff --git a/check/commands/phash/drop.py b/check/commands/phash/drop.py
new file mode 100644
index 0000000..40a8261
--- /dev/null
+++ b/check/commands/phash/drop.py
@@ -0,0 +1,22 @@
+"""
+Drop the database (useful when testing)
+"""
+
+import click
+import glob
+
+from app.models.sql_factory import Base, engine
+
+@click.command()
+@click.option('-f', '--force', 'opt_force', is_flag=True,
+  help='Actually drop the database')
+@click.pass_context
+def cli(ctx, opt_force):
+  """
+  Drop the database
+  """
+  if opt_force:
+    print('Dropping the database...!')
+    Base.metadata.drop_all(engine)
+  else:
+    print('Will foolishly drop the database if the --force flag is passed')
\ No newline at end of file
diff --git a/check/commands/phash/load.py b/check/commands/phash/load.py
new file mode 100644
index 0000000..ce9ceef
--- /dev/null
+++ b/check/commands/phash/load.py
@@ -0,0 +1,23 @@
+"""
+Loop over a directory of images
+ - Compute their phashes
+ - Optionally upload them to s3?
+"""
+
+import click
+import glob
+
+from app.models.sql_factory import add_phash_by_filename
+
+@click.command()
+@click.option('-i', '--input', 'opt_input_glob',
+  required=True,
+  help="File glob to add -- e.g. '../docs/images/*.jpg'")
+@click.pass_context
+def cli(ctx, opt_input_glob):
+  """
+  Add a directory of images
+  """
+  print('Adding a directory...')
+  for fn in glob.iglob(opt_input_glob):
+    add_phash_by_filename(fn)
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))
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))