1 files changed, 20 insertions, 14 deletions

diff --git a/megapixels/commands/demo/face_search.py b/megapixels/commands/demo/face_search.py
index 0452cc9d..ca0b8016 100644
--- a/megapixels/commands/demo/face_search.py
+++ b/megapixels/commands/demo/face_search.py
@@ -6,23 +6,27 @@ from app.utils import click_utils
 from app.settings import app_cfg as cfg
 from app.utils.logger_utils import Logger
 
+log = Logger.getLogger()
+
 @click.command()
 @click.option('-i', '--input', 'opt_fp_in', required=True,
-  help='Input face image')
+  help='File to lookup')
 @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))
+  help=click_utils.show_help(types.DataStore))
 @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.option('--results', 'opt_results', default=5, 
+  help='Number of match results to display')
+@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):
+def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_results, opt_gpu):
   """Display image info"""
   
   import sys
@@ -31,18 +35,22 @@ def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_gpu):
   from pathlib import Path
   import time
 
-  import imutils
   import pandas as pd
   import cv2 as cv
-  import dlib
   from tqdm import tqdm
-  
+  import imutils
+
   from app.utils import file_utils, im_utils
-  from app.models.data_store import DataStore, DataStoreS3
+  from app.models.data_store import DataStore
   from app.processors import face_detector
   from app.processors import face_recognition
 
   log = Logger.getLogger()
+  # init dataset
+  dataset = Dataset(opt_data_store, opt_dataset)
+  dataset.load_face_vectors()
+  dataset.load_records()
+  dataset.load_identities()
 
   # init face detection
   detector = face_detector.DetectorDLIBHOG()
@@ -52,6 +60,7 @@ def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, 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]
@@ -64,12 +73,9 @@ def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_gpu):
   
   # 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)
+  image_records = dataset.find_matches(vec_query, n_results=opt_results)
 
   # summary
   ims_match = [im_query]
@@ -91,4 +97,4 @@ def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_gpu):
       sys.exit()
     elif k != 255:
       # any key to continue
-      break  
\ No newline at end of file
+      break