adding 3D landmarks...

5 files changed, 76 insertions, 10 deletions

diff --git a/megapixels/app/processors/face_landmarks.py b/megapixels/app/processors/face_landmarks_2d.py
index dfcb9ee8..e8ce93c1 100644
--- a/megapixels/app/processors/face_landmarks.py
+++ b/megapixels/app/processors/face_landmarks_2d.py
@@ -11,6 +11,33 @@ from app.settings import app_cfg as cfg
 from app.settings import types
 from app.models.bbox import BBox
 
+class LandmarksFaceAlignment:
+
+  # Estimates 2D facial landmarks
+  import face_alignment
+
+  def __init__(self, gpu=0):
+    self.log = logger_utils.Logger.getLogger()
+    device = f'cuda:{gpu}' if gpu > -1 else 'cpu'
+    self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, device=device, flip_input=True)
+
+  def landmarks(self, im, as_type=str):
+    '''Calculates the 3D facial landmarks
+    :param im: (numpy.ndarray) image
+    :param as_type: (str) or (list) type to return data
+    '''
+    preds = self.fa.get_landmarks(im)
+    # convert to comma separated ints
+    # storing data as "[1,2], [3,4]" is larger file size than storing as "1,2,3,4"
+    # storing a list object in Pandas seems to result in 30% larger CSV files
+    # TODO optimize this
+    preds_int = [list(map(int, x)) for x in preds[0]]  # list of ints
+    if as_type is str:
+      return ','.join([','.join(list(map(str,[x,y]))) for x,y in preds_int])
+    else:
+      return preds_int
+
+
 class LandmarksDLIB:
 
   def __init__(self):
@@ -39,7 +66,7 @@ class LandmarksMTCNN:
     from mtcnn.mtcnn import MTCNN
     self.detector = MTCNN()
 
-  def detect(self, im, opt_size=None, opt_conf_thresh=None, opt_pyramids=None):
+  def landmarks(self, im, opt_size=None, opt_conf_thresh=None, opt_pyramids=None):
     '''Detects face using MTCNN and returns (list) of BBox
     :param im: (numpy.ndarray) image
     :returns list of BBox
diff --git a/megapixels/app/processors/face_landmarks_3d.py b/megapixels/app/processors/face_landmarks_3d.py
index 28aff592..3663364c 100644
--- a/megapixels/app/processors/face_landmarks_3d.py
+++ b/megapixels/app/processors/face_landmarks_3d.py
@@ -13,24 +13,63 @@ from app.settings import app_cfg as cfg
 from app.settings import types
 
 
+class FaceLandmarks2D:
+
+  # Estimates 2D facial landmarks
+  import face_alignment
+
+  def __init__(self, gpu=0):
+    self.log = logger_utils.Logger.getLogger()
+    device = f'cuda:{gpu}' if gpu > -1 else 'cpu'
+    self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, device=device, flip_input=True)
+
+  def landmarks(self, im, as_type=str):
+    '''Calculates the 3D facial landmarks
+    :param im: (numpy.ndarray) image
+    :param as_type: (str) or (list) type to return data
+    '''
+    preds = self.fa.get_landmarks(im)
+    # convert to comma separated ints
+    # storing data as "[1,2], [3,4]" is larger file size than storing as "1,2,3,4"
+    # storing a list object in Pandas seems to result in 30% larger CSV files
+    # TODO optimize this
+    preds_int = [list(map(int, x)) for x in preds[0]]  # list of ints
+    if as_type is str:
+      return ','.join([','.join(list(map(str,[x,y]))) for x,y in preds_int])
+    else
+      return preds_int
+
 
 class FaceLandmarks3D:
 
   # Estimates 3D facial landmarks
   import face_alignment
-  from skimage import io
 
-  def __init__(self):
+  def __init__(self, gpu=0):
     self.log = logger_utils.Logger.getLogger()
-    self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)
+    device = f'cuda:{gpu}' if gpu > -1 else 'cpu'
+    self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._3D, device=device, flip_input=False)
 
-  def landmarks(self, im):
+  def landmarks(self, im, as_type=str):
+    '''Calculates the 3D facial landmarks
+    :param im: (numpy.ndarray) image
+    :param as_type: (str) or (list) type to return data
+    '''
     preds = self.fa.get_landmarks(im)
-    return preds
+    # convert to comma separated ints
+    # storing data as "[1,2], [3,4]" is larger file size than storing as "1,2,3,4"
+    # storing a list object in Pandas seems to result in 30% larger CSV files
+    # TODO optimize this
+    preds_int = [list(map(int, x)) for x in preds[0]]  # list of ints
+    if as_type is str:
+      return ','.join([','.join(list(map(str,[x,y]))) for x,y in preds_int])
+    else
+      return preds_int
 
   def draw(self, im):
     '''draws landmarks in 3d scene'''
 
+    # TODO
     '''
     import face_alignment
     import numpy as np
@@ -74,4 +113,4 @@ class FaceLandmarks3D:
     ax.set_xlim(ax.get_xlim()[::-1])
     plt.show()
     '''
-    return False
\ No newline at end of file
+    return im
\ No newline at end of file
diff --git a/megapixels/app/settings/types.py b/megapixels/app/settings/types.py
index 0805c5bd..c2e2caf7 100644
--- a/megapixels/app/settings/types.py
+++ b/megapixels/app/settings/types.py
@@ -45,8 +45,8 @@ class LogLevel(Enum):
 # --------------------------------------------------------------------
 
 class Metadata(Enum):
-  IDENTITY, FILE_RECORD, FACE_VECTOR, FACE_POSE, FACE_ROI, FACE_LANDMARKS_68, \
-    FACE_LANDMARKS_3D = range(7)
+  IDENTITY, FILE_RECORD, FACE_VECTOR, FACE_POSE, FACE_ROI, FACE_LANDMARKS_2D_68, \
+    FACE_LANDMARKS_3D_68 = range(7)
 
 class Dataset(Enum):
   LFW, VGG_FACE2, MSCELEB, UCCS, UMD_FACES = range(5)
diff --git a/megapixels/commands/cv/face_landmarks_3d.py b/megapixels/commands/cv/face_landmark.py
index 03ef8fc2..03ef8fc2 100644
--- a/megapixels/commands/cv/face_landmarks_3d.py
+++ b/megapixels/commands/cv/face_landmark.py
diff --git a/megapixels/commands/cv/face_pose.py b/megapixels/commands/cv/face_pose.py
index 9979ad34..4e35210c 100644
--- a/megapixels/commands/cv/face_pose.py
+++ b/megapixels/commands/cv/face_pose.py
@@ -57,7 +57,7 @@ def cli(ctx, opt_fp_in, opt_fp_out, opt_dir_media, opt_data_store, opt_dataset,
 
   from app.models.bbox import BBox
   from app.utils import logger_utils, file_utils, im_utils
-  from app.processors.face_landmarks import LandmarksDLIB
+  from app.processors.face_landmarks_2d import LandmarksDLIB
   from app.processors.face_pose import FacePoseDLIB
   from app.models.data_store import DataStore