add demos, in progress

4 files changed, 23 insertions, 10 deletions

diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py
index 6bf27576..c0762564 100644
--- a/megapixels/app/processors/face_detector.py
+++ b/megapixels/app/processors/face_detector.py
@@ -1,3 +1,4 @@
+import sys
 import os
 from os.path import join
 from pathlib import Path
@@ -30,8 +31,6 @@ class DetectorMTCNN:
     :returns list of BBox
     '''
     bboxes = []
-    #conf_thresh = self.conf_thresh if conf_thresh is None else conf_thresh
-    #pyramids = self.pyramids if pyramids is None else pyramids
     dnn_size = self.dnn_size if size is None else size
 
     im = im_utils.resize(im, width=dnn_size[0], height=dnn_size[1])
@@ -72,6 +71,9 @@ class DetectorDLIBCNN:
     import dlib
     self.log = logger_utils.Logger.getLogger()
     cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', '')
+    if dlib.DLIB_USE_CUDA and gpu < 0:
+      self.log.error('dlib was compiled with CUDA but you selected CPU. Use GPU >= 0 if dlib.DLIB_USE_CUDA')
+      sys.exit()
     os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
     self.log.info('load model: {}'.format(cfg.DIR_MODELS_DLIB_CNN))
     self.detector = dlib.cnn_face_detection_model_v1(cfg.DIR_MODELS_DLIB_CNN)
diff --git a/megapixels/app/processors/face_landmarks.py b/megapixels/app/processors/face_landmarks.py
index 8086ba1e..171fc666 100644
--- a/megapixels/app/processors/face_landmarks.py
+++ b/megapixels/app/processors/face_landmarks.py
@@ -83,8 +83,11 @@ class Dlib2D(Landmarks2D):
     self.log.info(f'loaded predictor model: {model}')
 
   def landmarks(self, im, bbox):
-    # Draw high-confidence faces
-    dim_wh = im.shape[:2][::-1]
+    '''Generates 68-pt landmarks using dlib predictor
+    :param im: (numpy.ndarray) BGR image
+    :param bbox: (app.models.BBox) dimensioned
+    :returns (list) of (int, int) for x,y values
+    '''
     bbox = bbox.to_dlib()
     im_gray = cv.cvtColor(im, cv.COLOR_BGR2GRAY)
     points = [[p.x, p.y] for p in self.predictor(im_gray, bbox).parts()]
@@ -168,8 +171,8 @@ class Landmarks3D:
         points_formatted[f'{d}{idx}'] = pt[j]
     return points_formatted
 
-  def normalize(self, points, dim):
-    return [np.array(p)/dim for p in points]  # divides each point by w,h dim
+  # def normalize(self, points):
+  #   '''TODO'''
 
     
 class FaceAlignment3D_68(Landmarks3D):
@@ -182,13 +185,14 @@ class FaceAlignment3D_68(Landmarks3D):
     device = f'cuda:{gpu}' if gpu > -1 else 'cpu'
     self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._3D, device=device, flip_input=flip_input)
 
-  def landmarks(self, im, as_type=str):
+  def landmarks(self, im, rect):
     '''Calculates the 3D facial landmarks
     :param im: (numpy.ndarray) BGR image
+    :param rect: (list) of face (x1, y1, x2, y2)
     :returns (list) of 68 (int) (tuples) as (x,y, z) 
     '''
     # predict landmarks
-    points = self.fa.get_landmarks(im)  # returns array of arrays of 68 3D pts/face
+    points = self.fa.get_landmarks(im, [rect])  # returns array of arrays of 68 3D pts/face
     # convert to data type
     points = [list(map(int, p)) for p in points[0]]
     return points
\ No newline at end of file
diff --git a/megapixels/app/processors/face_landmarks_3d.py b/megapixels/app/processors/face_landmarks_3d.py
index 470d263c..5a0d6097 100644
--- a/megapixels/app/processors/face_landmarks_3d.py
+++ b/megapixels/app/processors/face_landmarks_3d.py
@@ -26,14 +26,15 @@ class FaceAlignment3D(Landmarks3D):
   # Estimates 3D facial landmarks
   import face_alignment
 
-  def __init__(self, gpu=0, flip_input=False):
+  def __init__(self, gpu=0, flip_input=True):
     super().__init__()
     device = f'cuda:{gpu}' if gpu > -1 else 'cpu'
     self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._3D, device=device, flip_input=flip_input)
 
-  def landmarks(self, im, as_type=str):
+  def landmarks(self, im, bbox, as_type=str):
     '''Calculates the 3D facial landmarks
     :param im: (numpy.ndarray) image
+    :param bbox: (BBox) dimensioned to real (int) sizes
     :param as_type: (str) or (list) type to return data
     '''
     preds = self.fa.get_landmarks(im)
diff --git a/megapixels/app/processors/face_pose.py b/megapixels/app/processors/face_pose.py
index 8bc95f8d..5ac510ec 100644
--- a/megapixels/app/processors/face_pose.py
+++ b/megapixels/app/processors/face_pose.py
@@ -25,6 +25,12 @@ class FacePoseDLIB:
 
 
   def pose(self, landmarks, dim):
+    '''Returns face pose information
+    :param landmarks: (list) of 68 (int, int) xy tuples
+    :param dim: (tuple|list) of image (width, height)
+    :returns (dict) of pose attributes
+    '''
+
     # computes pose using 6 / 68 points from dlib face landmarks
     # based on learnopencv.com and 
     # https://github.com/jerryhouuu/Face-Yaw-Roll-Pitch-from-Pose-Estimation-using-OpenCV/