add 3d render

1 files changed, 95 insertions, 2 deletions

diff --git a/megapixels/app/utils/draw_utils.py b/megapixels/app/utils/draw_utils.py
index 7083c956..3a389e68 100644
--- a/megapixels/app/utils/draw_utils.py
+++ b/megapixels/app/utils/draw_utils.py
@@ -1,8 +1,102 @@
 import sys
+from math import sqrt
 
+import numpy as np
 import cv2 as cv
 
 
+end_list = np.array([17, 22, 27, 42, 48, 31, 36, 68], dtype=np.int32) - 1
+
+# ---------------------------------------------------------------------------
+#
+# 3D landmark drawing utilities
+#
+# ---------------------------------------------------------------------------
+
+def plot_keypoints(im, kpts):
+  '''Draw 68 key points
+  :param im: the input im
+  :param kpts: (68, 3). flattened list
+  '''
+  im = im.copy()
+  kpts = np.round(kpts).astype(np.int32)
+  for i in range(kpts.shape[0]):
+    st = kpts[i, :2]
+    im = cv.circle(im, (st[0], st[1]), 1, (0, 0, 255), 2)
+    if i in end_list:
+      continue
+    ed = kpts[i + 1, :2]
+    im = cv.line(im, (st[0], st[1]), (ed[0], ed[1]), (255, 255, 255), 1)
+  return im
+
+
+def calc_hypotenuse(pts):
+  bbox = [min(pts[0, :]), min(pts[1, :]), max(pts[0, :]), max(pts[1, :])]
+  center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
+  radius = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
+  bbox = [center[0] - radius, center[1] - radius, center[0] + radius, center[1] + radius]
+  llength = sqrt((bbox[2] - bbox[0]) ** 2 + (bbox[3] - bbox[1]) ** 2)
+  return llength / 3
+
+def build_camera_box(rear_size=90):
+  point_3d = []
+  rear_depth = 0
+  point_3d.append((-rear_size, -rear_size, rear_depth))
+  point_3d.append((-rear_size, rear_size, rear_depth))
+  point_3d.append((rear_size, rear_size, rear_depth))
+  point_3d.append((rear_size, -rear_size, rear_depth))
+  point_3d.append((-rear_size, -rear_size, rear_depth))
+
+  front_size = int(4 / 3 * rear_size)
+  front_depth = int(4 / 3 * rear_size)
+  point_3d.append((-front_size, -front_size, front_depth))
+  point_3d.append((-front_size, front_size, front_depth))
+  point_3d.append((front_size, front_size, front_depth))
+  point_3d.append((front_size, -front_size, front_depth))
+  point_3d.append((-front_size, -front_size, front_depth))
+  point_3d = np.array(point_3d, dtype=np.float).reshape(-1, 3)
+
+  return point_3d
+
+
+def plot_pose_box(im, Ps, pts68s, color=(40, 255, 0), line_width=2):
+  '''Draw a 3D box as annotation of pose. 
+    ref: https://github.com/yinguobing/head-pose-estimation/blob/master/pose_estimator.py
+  :param image: the input image
+  :param P: (3, 4). Affine Camera Matrix.
+  :param kpts: (2, 68) or (3, 68)
+  '''
+  im_draw = im.copy()
+  if not isinstance(pts68s, list):
+    pts68s = [pts68s]
+  
+  if not isinstance(Ps, list):
+    Ps = [Ps]
+  
+  for i in range(len(pts68s)):
+    pts68 = pts68s[i]
+    llength = calc_hypotenuse(pts68)
+    point_3d = build_camera_box(llength)
+    P = Ps[i]
+
+    # Map to 2d im points
+    point_3d_homo = np.hstack((point_3d, np.ones([point_3d.shape[0], 1])))  # n x 4
+    point_2d = point_3d_homo.dot(P.T)[:, :2]
+
+    point_2d[:, 1] = - point_2d[:, 1]
+    point_2d[:, :2] = point_2d[:, :2] - np.mean(point_2d[:4, :2], 0) + np.mean(pts68[:2, :27], 1)
+    point_2d = np.int32(point_2d.reshape(-1, 2))
+
+    # Draw all the lines
+    cv.polylines(im_draw, [point_2d], True, color, line_width, cv.LINE_AA)
+    cv.line(im_draw, tuple(point_2d[1]), tuple(point_2d[6]), color, line_width, cv.LINE_AA)
+    cv.line(im_draw, tuple(point_2d[2]), tuple(point_2d[7]), color, line_width, cv.LINE_AA)
+    cv.line(im_draw, tuple(point_2d[3]), tuple(point_2d[8]), color, line_width, cv.LINE_AA)
+
+    return im_draw
+
+
+
 # ---------------------------------------------------------------------------
 #
 # OpenCV drawing functions
@@ -11,7 +105,6 @@ import cv2 as cv
 
 pose_types = {'pitch': (0,0,255), 'roll': (255,0,0), 'yaw': (0,255,0)}
 
-
 def draw_landmarks2D(im, points, radius=3, color=(0,255,0), stroke_weight=2):
   '''Draws facial landmarks, either 5pt or 68pt
   '''
@@ -27,7 +120,7 @@ def draw_landmarks3D(im, points, radius=3, color=(0,255,0), stroke_weight=2):
 
 
 def draw_bbox(im, bbox, color=(0,255,0), stroke_weight=2):
-  '''Draws a dimensioned (not-normalized) BBox onto cv2 image
+  '''Draws a dimensioned (not-normalized) BBox onto cv image
   '''
   cv.rectangle(im, bbox.pt_tl, bbox.pt_br, color, stroke_weight)