change bbox to norm, refine face extractor

1 files changed, 19 insertions, 21 deletions

diff --git a/megapixels/commands/demo/face_search.py b/megapixels/commands/demo/face_search.py
index ca0b8016..d50f5c73 100644
--- a/megapixels/commands/demo/face_search.py
+++ b/megapixels/commands/demo/face_search.py
@@ -13,7 +13,7 @@ log = Logger.getLogger()
   help='File to lookup')
 @click.option('--data_store', 'opt_data_store',
   type=cfg.DataStoreVar,
-  default=click_utils.get_default(types.DataStore.SSD),
+  default=click_utils.get_default(types.DataStore.HDD),
   show_default=True,
   help=click_utils.show_help(types.DataStore))
 @click.option('--dataset', 'opt_dataset',
@@ -40,44 +40,47 @@ def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_results, opt_gpu):
   from tqdm import tqdm
   import imutils
 
-  from app.utils import file_utils, im_utils
+  from app.utils import file_utils, im_utils, display_utils, draw_utils
   from app.models.data_store import DataStore
   from app.processors import face_detector
-  from app.processors import face_recognition
+  from app.processors import face_extractor
 
   log = Logger.getLogger()
+
   # init dataset
   dataset = Dataset(opt_data_store, opt_dataset)
   dataset.load_face_vectors()
   dataset.load_records()
-  dataset.load_identities()
+  # dataset.load_identities()
 
   # init face detection
-  detector = face_detector.DetectorDLIBHOG()
+  detector = face_detector.DetectorCVDNN()
 
-  # init face recognition
-  recognition = face_recognition.RecognitionDLIB(gpu=opt_gpu)
+  # init face extractor
+  extractor = face_extractor.ExtractorVGG()
   
   # 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]
+  bbox_norm = bboxes[0]
   dim = im_query.shape[:2][::-1]
-  bbox = bbox.to_dim(dim)  # convert back to real dimensions
+  bbox_dim = bbox_norm.to_dim(dim)  # convert back to real dimensions
 
-  if not bbox:
+  if not bbox_norm:
     log.error('No face detected. Exiting')
     return
   
   # extract the face vectors
-  vec_query = recognition.vec(im_query, bbox)
+  vec_query = extractor.extract(im_query, bbox_norm)
+  log.debug(f'len query: {len(vec_query)}')
 
   # find matches
   image_records = dataset.find_matches(vec_query, n_results=opt_results)
 
   # summary
+  im_query = draw_utils.draw_bbox(im_query, bbox_norm, stroke_weight=8)
   ims_match = [im_query]
   for image_record in image_records:
     image_record.summarize()
@@ -85,16 +88,11 @@ def cli(ctx, opt_fp_in, opt_data_store, opt_dataset, opt_results, opt_gpu):
     im_match = cv.imread(image_record.filepath)
     ims_match.append(im_match)
 
+  # make montages of most similar faces
   montages = imutils.build_montages(ims_match, (256, 256), (3,2))
 
+  # display 
   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  
+    cv.imshow(f'{opt_dataset.name.upper()}: page {i}', montage)
+  
+  display_utils.handle_keyboard()