1 files changed, 45 insertions, 6 deletions
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