1 files changed, 10 insertions, 28 deletions
diff --git a/megapixels/app/processors/face_landmarks_3d.py b/megapixels/app/processors/face_landmarks_3d.py
index 3663364c..470d263c 100644
--- a/megapixels/app/processors/face_landmarks_3d.py
+++ b/megapixels/app/processors/face_landmarks_3d.py
@@ -12,43 +12,24 @@ from app.models.bbox import BBox
 from app.settings import app_cfg as cfg
 from app.settings import types
 
+class Landmarks3D:
 
-class FaceLandmarks2D:
-
-  # Estimates 2D facial landmarks
-  import face_alignment
-
-  def __init__(self, gpu=0):
+  def __init__(self):
     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
 
+  def landmarks(self, im, bbox):
+    pass
 
-class FaceLandmarks3D:
+    
+class FaceAlignment3D(Landmarks3D):
 
   # Estimates 3D facial landmarks
   import face_alignment
 
-  def __init__(self, gpu=0):
-    self.log = logger_utils.Logger.getLogger()
+  def __init__(self, gpu=0, flip_input=False):
+    super().__init__()
     device = f'cuda:{gpu}' if gpu > -1 else 'cpu'
-    self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._3D, device=device, flip_input=False)
+    self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._3D, device=device, flip_input=flip_input)
 
   def landmarks(self, im, as_type=str):
     '''Calculates the 3D facial landmarks
@@ -66,6 +47,7 @@ class FaceLandmarks3D:
     else
       return preds_int
 
+
   def draw(self, im):
     '''draws landmarks in 3d scene'''