diff options
Diffstat (limited to 'megapixels/commands/demo/face_search.py')
| -rw-r--r-- | megapixels/commands/demo/face_search.py | 26 |
1 files changed, 15 insertions, 11 deletions
diff --git a/megapixels/commands/demo/face_search.py b/megapixels/commands/demo/face_search.py index 0452cc9d..34a25762 100644 --- a/megapixels/commands/demo/face_search.py +++ b/megapixels/commands/demo/face_search.py @@ -6,9 +6,11 @@ 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), @@ -19,8 +21,8 @@ from app.utils.logger_utils import Logger 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('--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""" @@ -31,18 +33,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 +58,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,9 +71,6 @@ 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) @@ -91,4 +95,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 |