add vgg

1 files changed, 24 insertions, 12 deletions

diff --git a/megapixels/commands/cv/face_vector.py b/megapixels/commands/cv/face_vector.py
index 4df647f5..9e9f6396 100644
--- a/megapixels/commands/cv/face_vector.py
+++ b/megapixels/commands/cv/face_vector.py
@@ -1,5 +1,8 @@
 """
 Converts ROIs to face vector
+NB: the VGG Face2 extractor should be used with MTCNN ROIs (not square)
+  the DLIB face extractor should be used with DLIB ROIs (square)
+see https://github.com/ox-vgg/vgg_face2 for TAR@FAR
 """
 
 import click
@@ -26,10 +29,14 @@ from app.settings import app_cfg as cfg
 @click.option('--size', 'opt_size', 
   type=(int, int), default=(300, 300),
   help='Output image size')
+@click.option('-e', '--extractor', 'opt_extractor',
+  default=types.FaceExtractor.VGG,
+  type=cfg.FaceExtractorVar,
+  help='Type of extractor framework/network to use')
 @click.option('-j', '--jitters', 'opt_jitters', default=cfg.DLIB_FACEREC_JITTERS,
-  help='Number of jitters')
-@click.option('-p', '--padding', 'opt_padding', default=cfg.DLIB_FACEREC_PADDING,
-  help='Percentage padding')
+  help='Number of jitters (only for dlib')
+@click.option('-p', '--padding', 'opt_padding', default=cfg.FACEREC_PADDING,
+  help='Percentage ROI padding')
 @click.option('--slice', 'opt_slice', type=(int, int), default=(None, None),
   help='Slice list of files')
 @click.option('-f', '--force', 'opt_force', is_flag=True,
@@ -38,7 +45,7 @@ from app.settings import app_cfg as cfg
   help='GPU index')
 @click.pass_context
 def cli(ctx, opt_fp_out, opt_dir_media, opt_data_store, opt_dataset, opt_size, 
-  opt_slice, opt_force, opt_gpu, opt_jitters, opt_padding):
+  opt_extractor, opt_slice, opt_force, opt_gpu, opt_jitters, opt_padding):
   """Converts face ROIs to vectors"""
   
   import sys
@@ -56,7 +63,7 @@ def cli(ctx, opt_fp_out, opt_dir_media, opt_data_store, opt_dataset, opt_size,
   from app.models.bbox import BBox
   from app.models.data_store import DataStore
   from app.utils import logger_utils, file_utils, im_utils
-  from app.processors import face_recognition
+  from app.processors import face_extractor
 
 
   # -------------------------------------------------
@@ -73,7 +80,11 @@ def cli(ctx, opt_fp_out, opt_dir_media, opt_data_store, opt_dataset, opt_size,
     return
 
   # init face processors
-  facerec = face_recognition.RecognitionDLIB()
+  if opt_extractor == types.FaceExtractor.DLIB:
+    log.debug('set dlib')
+    extractor = face_extractor.ExtractorDLIB(gpu=opt_gpu, jitters=opt_jitters)
+  elif opt_extractor == types.FaceExtractor.VGG:
+    extractor = face_extractor.ExtractorVGG()
 
   # load data
   fp_record = data_store.metadata(types.Metadata.FILE_RECORD)
@@ -85,7 +96,8 @@ def cli(ctx, opt_fp_out, opt_dir_media, opt_data_store, opt_dataset, opt_size,
     df_roi = df_roi[opt_slice[0]:opt_slice[1]]
   
   # -------------------------------------------------
-  # process here
+  # process images
+
   df_img_groups = df_roi.groupby('record_index')
   log.debug('processing {:,} groups'.format(len(df_img_groups)))
 
@@ -99,17 +111,17 @@ def cli(ctx, opt_fp_out, opt_dir_media, opt_data_store, opt_dataset, opt_size,
       # get bbox
       x, y, w, h = df_img.x, df_img.y, df_img.w, df_img.h
       dim = (ds_record.width, ds_record.height)
-      #dim = im.shape[:2][::-1]
       # get face vector
-      bbox_dim = BBox.from_xywh(x, y, w, h).to_dim(dim)  # convert to int real dimensions
+      bbox = BBox.from_xywh(x, y, w, h)  # norm
       # compute vec
-      # padding=opt_padding not yet implemented in dlib===19.16 but merged in master
-      vec = facerec.vec(im, bbox_dim, jitters=opt_jitters)
-      vec_flat = facerec.flatten(vec)
+      vec = extractor.extract(im, bbox)  # use normalized BBox
+      vec_flat = extractor.flatten(vec)
       vec_flat['roi_index'] = roi_index
       vec_flat['record_index'] = record_index
       vecs.append(vec_flat)
 
+  # -------------------------------------------------
+  # save data
 
   # create DataFrame and save to CSV
   df = pd.DataFrame.from_dict(vecs)