summaryrefslogtreecommitdiff
path: root/megapixels/commands/demo/face_vector.py
diff options
context:
space:
mode:
Diffstat (limited to 'megapixels/commands/demo/face_vector.py')
-rw-r--r--megapixels/commands/demo/face_vector.py79
1 files changed, 79 insertions, 0 deletions
diff --git a/megapixels/commands/demo/face_vector.py b/megapixels/commands/demo/face_vector.py
new file mode 100644
index 00000000..1104f923
--- /dev/null
+++ b/megapixels/commands/demo/face_vector.py
@@ -0,0 +1,79 @@
+"""
+Crop images to prepare for training
+"""
+
+import click
+# from PIL import Image, ImageOps, ImageFilter, ImageDraw
+
+from app.settings import types
+from app.utils import click_utils
+from app.settings import app_cfg as cfg
+
+
+@click.command()
+@click.option('-i', '--input', 'opt_fp_in', default=None, required=True,
+ help='Image filepath')
+@click.option('--size', 'opt_size',
+ type=(int, int), default=(300, 300),
+ help='Output image size')
+@click.option('-g', '--gpu', 'opt_gpu', default=0,
+ help='GPU index')
+@click.option('--display/--no-display', 'opt_display', is_flag=True, default=False,
+ help='Display detections to debug')
+@click.pass_context
+def cli(ctx, opt_fp_in, opt_gpu, opt_size, opt_display):
+ """Demo generating face vector"""
+
+ import sys
+ import os
+ from os.path import join
+ from pathlib import Path
+ import time
+
+ import numpy as np
+ import cv2 as cv
+ import dlib # NB: keep a reference in main file if using dlib detector processors
+
+ from app.utils import logger_utils, file_utils, im_utils, display_utils, draw_utils
+ from app.utils import plot_utils
+ from app.processors import face_detector
+ from app.models.data_store import DataStore
+
+ # -------------------------------------------------
+ # init here
+
+ log = logger_utils.Logger.getLogger()
+
+ # ----------------------------------------------------------------------------
+ # load image
+
+ im = cv.imread(opt_fp_in)
+ im_resized = im_utils.resize(im, width=opt_size[0], height=opt_size[1])
+
+
+ # ----------------------------------------------------------------------------
+ # detect face
+
+ face_detector = face_detector.DetectorDLIBCNN(gpu=opt_gpu) # -1 for CPU
+ bboxes = face_detector.detect(im_resized, largest=True)
+ bbox = bboxes[0]
+ dim = im_resized.shape[:2][::-1]
+ bbox_dim = bbox.to_dim(dim)
+ if not bbox:
+ log.error('no face detected')
+ return
+
+
+ # ----------------------------------------------------------------------------
+ # generate face vectors, only to test if feature extraction works
+
+ from app.processors import face_recognition
+ face_rec = face_recognition.RecognitionDLIB()
+ vec = face_rec.vec(im_resized, bbox_dim)
+ log.info(f'generated vector. showing vec[0:10]:')
+ log.info(f'\n{vec[0:10]}')
+
+ if opt_display:
+ draw_utils.draw_bbox(im_resized, bbox_dim)
+ cv.imshow('Original', im_resized)
+ display_utils.handle_keyboard() \ No newline at end of file