cli suite working

5 files changed, 152 insertions, 0 deletions

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))