fix roi index, clean up pose, roi, records, vector

4 files changed, 39 insertions, 19 deletions

diff --git a/megapixels/app/models/data_store.py b/megapixels/app/models/data_store.py
index 7b6bef21..626c9da4 100644
--- a/megapixels/app/models/data_store.py
+++ b/megapixels/app/models/data_store.py
@@ -28,6 +28,8 @@ class DataStore:
     return join(self.dir_media, 'original')
 
   def face(self, subdir, fn, ext):
+    if subdir == '' or subdir is None:
+      subdir = '.'
     return join(self.dir_media, 'original', subdir, f'{fn}.{ext}')
 
   def face_crop(self, subdir, fn, ext):
diff --git a/megapixels/app/processors/face_detector.py b/megapixels/app/processors/face_detector.py
index 3a90c557..75ba54d4 100644
--- a/megapixels/app/processors/face_detector.py
+++ b/megapixels/app/processors/face_detector.py
@@ -24,7 +24,7 @@ class DetectorMTCNN:
     from mtcnn.mtcnn import MTCNN
     self.detector = MTCNN()
 
-  def detect(self, im, size=(400,400), conf_thresh=None, pyramids=None, largest=False):
+  def detect(self, im, size=(400,400), conf_thresh=None, pyramids=None, largest=False, zone=None):
     '''Detects face using MTCNN and returns (list) of BBox
     :param im: (numpy.ndarray) image
     :returns list of BBox
@@ -79,7 +79,7 @@ class DetectorDLIBCNN:
     self.detector = dlib.cnn_face_detection_model_v1(cfg.DIR_MODELS_DLIB_CNN)
     os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devices  # reset
 
-  def detect(self, im, size=None, conf_thresh=None, pyramids=None, largest=False):
+  def detect(self, im, size=None, conf_thresh=None, pyramids=None, largest=False, zone=None):
     bboxes = []
     conf_thresh = self.conf_thresh if conf_thresh is None else conf_thresh
     pyramids = self.pyramids if pyramids is None else pyramids
@@ -96,6 +96,10 @@ class DetectorDLIBCNN:
         bbox = BBox.from_dlib_dim(mmod_rect.rect, dim)
         bboxes.append(bbox)
 
+    if zone:
+      bboxes = [b for b in bboxes if b.cx > zone[0] and b.cx < 1.0 - zone[0] \
+        and b.cy > zone[1] and b.cy < 1.0 - zone[1]]
+
     if largest and len(bboxes) > 1:
       # only keep largest
       bboxes.sort(key=operator.attrgetter('area'), reverse=True)
diff --git a/megapixels/app/processors/face_pose.py b/megapixels/app/processors/face_pose.py
index f2548b32..96281637 100644
--- a/megapixels/app/processors/face_pose.py
+++ b/megapixels/app/processors/face_pose.py
@@ -18,12 +18,13 @@ class FacePoseDLIB:
 
 
   dnn_size = (400, 400)
+  pose_types = {'pitch': (0,0,255), 'roll': (255,0,0), 'yaw': (0,255,0)}
 
   def __init__(self):
     pass
 
 
-  def pose(self, landmarks, dim, project_points=False):
+  def pose(self, landmarks, dim):
     # 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/
@@ -68,12 +69,17 @@ class FacePoseDLIB:
     result = {}
 
     # project points
-    if project_points:
-      pts_im, jac = cv.projectPoints(axis, rot_vec, tran_vec, cam_mat, dist_coeffs)
-      pts_model, jac2 = cv.projectPoints(model_points, rot_vec, tran_vec, cam_mat, dist_coeffs)
-      result['points_model'] = pts_model
-      result['points_image'] = pts_im
-      result['point_nose'] = tuple(landmarks[pose_points_idx[0]])
+    #if project_points:
+    pts_im, jac = cv.projectPoints(axis, rot_vec, tran_vec, cam_mat, dist_coeffs)
+    pts_model, jac2 = cv.projectPoints(model_points, rot_vec, tran_vec, cam_mat, dist_coeffs)
+    #result['points_model'] = pts_model
+    #result['points_image'] = pts_im
+    result['points'] = {
+      'pitch': pts_im[0],
+      'roll': pts_im[2],
+      'yaw': pts_im[1]
+      }
+    result['point_nose'] = tuple(landmarks[pose_points_idx[0]])
 
     rvec_matrix = cv.Rodrigues(rot_vec)[0]
     
@@ -84,21 +90,23 @@ class FacePoseDLIB:
     pitch = math.degrees(math.asin(math.sin(pitch)))
     roll = -math.degrees(math.asin(math.sin(roll)))
     yaw = math.degrees(math.asin(math.sin(yaw)))
-    degrees = {'pitch': pitch, 'roll': roll, 'yaw': yaw}
-    result['degrees'] = degrees
+    result['pitch'] = pitch
+    result['roll'] = roll
+    result['yaw'] = yaw
 
     return result
 
 
-  def draw_pose(self, im, pts_im, pts_model, pt_nose):
-    cv.line(im, pt_nose, tuple(pts_im[1].ravel()), (0,255,0), 3) #GREEN
-    cv.line(im, pt_nose, tuple(pts_im[0].ravel()), (255,0,), 3) #BLUE
-    cv.line(im, pt_nose, tuple(pts_im[2].ravel()), (0,0,255), 3) #RED
+  def draw_pose(self, im, pt_nose, image_pts):
+    cv.line(im, pt_nose, tuple(image_pts['pitch'].ravel()), self.pose_types['pitch'], 3)
+    cv.line(im, pt_nose, tuple(image_pts['yaw'].ravel()), self.pose_types['yaw'], 3)
+    cv.line(im, pt_nose, tuple(image_pts['roll'].ravel()), self.pose_types['roll'], 3)
 
 
-  def draw_degrees(self, im, degrees, color=(0,255,0)):
-    for i, item in enumerate(degrees.items()):
-      k, v = item
+  def draw_degrees(self, im, pose_data, color=(0,255,0)):
+    for i, pose_type in enumerate(self.pose_types.items()):
+      k, clr = pose_type
+      v = pose_data[k]
       t = '{}: {:.2f}'.format(k, v)
       origin = (10, 30 + (25 * i))
-      cv.putText(im, t, origin, cv.FONT_HERSHEY_SIMPLEX, 0.5, color, thickness=2, lineType=2)
\ No newline at end of file
+      cv.putText(im, t, origin, cv.FONT_HERSHEY_SIMPLEX, 0.5, clr, thickness=2, lineType=2)
\ No newline at end of file
diff --git a/megapixels/app/utils/im_utils.py b/megapixels/app/utils/im_utils.py
index d5e92aa3..e882c67f 100644
--- a/megapixels/app/utils/im_utils.py
+++ b/megapixels/app/utils/im_utils.py
@@ -21,6 +21,12 @@ import datetime
 
 
 
+def num_channels(im):
+  '''Returns number of channels in numpy.ndarray image'''
+  if len(im.shape) > 2:
+    return im.shape[2]
+  else:
+    return 1
 
 def is_grayscale(im, threshold=5):
   """Returns True if image is grayscale