diff options
Diffstat (limited to 'megapixels/commands/demo/face_search.py')
| -rw-r--r-- | megapixels/commands/demo/face_search.py | 40 |
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() |