1 files changed, 94 insertions, 0 deletions
diff --git a/megapixels/commands/demo/face_search.py b/megapixels/commands/demo/face_search.py
new file mode 100644
index 00000000..0452cc9d
--- /dev/null
+++ b/megapixels/commands/demo/face_search.py
@@ -0,0 +1,94 @@
+import click
+
+from app.settings import types
+from app.models.dataset import Dataset
+from app.utils import click_utils
+from app.settings import app_cfg as cfg
+from app.utils.logger_utils import Logger
+
+@click.command()
+@click.option('-i', '--input', 'opt_fp_in', required=True,
+  help='Input face image')
+@click.option('--data_store', 'opt_data_store',
+  type=cfg.DataStoreVar,
+  default=click_utils.get_default(types.DataStore.SSD),
+  show_default=True,
+  help=click_utils.show_help(types.Dataset))
+@click.option('--dataset', 'opt_dataset',
+  type=cfg.DatasetVar,
+  required=True,
+  show_default=True,
+  help=click_utils.show_help(types.Dataset))
+@click.option('--gpu', 'opt_gpu', default=0,
+  help='GPU index (use -1 for CPU)')
+@click.pass_context
+def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_gpu):
+  """Display image info"""
+  
+  import sys
+  from glob import glob
+  from os.path import join
+  from pathlib import Path
+  import time
+
+  import imutils
+  import pandas as pd
+  import cv2 as cv
+  import dlib
+  from tqdm import tqdm
+  
+  from app.utils import file_utils, im_utils
+  from app.models.data_store import DataStore, DataStoreS3
+  from app.processors import face_detector
+  from app.processors import face_recognition
+
+  log = Logger.getLogger()
+
+  # init face detection
+  detector = face_detector.DetectorDLIBHOG()
+
+  # init face recognition
+  recognition = face_recognition.RecognitionDLIB(gpu=opt_gpu)
+  
+  # load query image
+  im_query = cv.imread(opt_fp_in)
+  # get detection as BBox object
+  bboxes = detector.detect(im_query, largest=True)
+  bbox = bboxes[0]
+  dim = im_query.shape[:2][::-1]
+  bbox = bbox.to_dim(dim)  # convert back to real dimensions
+
+  if not bbox:
+    log.error('No face detected. Exiting')
+    return
+  
+  # extract the face vectors
+  vec_query = recognition.vec(im_query, bbox)
+  
+  # load dataset CSVs
+  dataset = Dataset(opt_data_store, opt_dataset)
+
+  # find matches
+  image_records = dataset.find_matches(vec_query, n_results=5)
+
+  # summary
+  ims_match = [im_query]
+  for image_record in image_records:
+    image_record.summarize()
+    log.info(f'{image_record.filepath}')
+    im_match = cv.imread(image_record.filepath)
+    ims_match.append(im_match)
+
+  montages = imutils.build_montages(ims_match, (256, 256), (3,2))
+
+  for i, montage in enumerate(montages):
+    cv.imshow(f'{i}', montage)
+  # cv gui
+  while True:
+    k = cv.waitKey(1) & 0xFF
+    if k == 27 or k == ord('q'):  # ESC
+      cv.destroyAllWindows()
+      sys.exit()
+    elif k != 255:
+      # any key to continue
+      break  
+\ No newline at end of file