externalize drawing, cleanup

2 files changed, 81 insertions, 0 deletions

diff --git a/megapixels/app/utils/display_utils.py b/megapixels/app/utils/display_utils.py
new file mode 100644
index 00000000..58e2feec
--- /dev/null
+++ b/megapixels/app/utils/display_utils.py
@@ -0,0 +1,16 @@
+import sys
+
+import cv2 as cv
+
+
+def handle_keyboard():
+  '''Used with cv.imshow('title', image) to wait for keyboard press
+  '''
+  while True:
+    k = cv.waitKey(1) & 0xFF
+    if k == 27 or k == ord('q'):  # ESC
+      cv.destroyAllWindows()
+      sys.exit()
+    elif k != 255:
+      # any key to continue
+      break
\ No newline at end of file
diff --git a/megapixels/app/utils/draw_utils.py b/megapixels/app/utils/draw_utils.py
new file mode 100644
index 00000000..f6d53609
--- /dev/null
+++ b/megapixels/app/utils/draw_utils.py
@@ -0,0 +1,65 @@
+import sys
+
+import cv2 as cv
+
+
+# ---------------------------------------------------------------------------
+#
+# OpenCV drawing functions
+#
+# ---------------------------------------------------------------------------
+
+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
+  '''
+  for x,y in points:
+    cv.circle(im, (x,y), radius, color, -1, cv.LINE_AA)
+
+
+def draw_landmarks3D(im, points, radius=3, color=(0,255,0), stroke_weight=2):
+  '''Draws 3D facial landmarks
+  '''
+  for x,y,z in points:
+    cv.circle(im, (x,y), radius, color, -1, cv.LINE_AA)
+
+
+def draw_bbox(im, bbox, color=(0,255,0), stroke_weight=2):
+  '''Draws a dimensioned (not-normalized) BBox onto cv2 image
+  '''
+  cv.rectangle(im, bbox.pt_tl, bbox.pt_br, color, stroke_weight)
+
+
+def draw_pose(im, pt_nose, image_pts):
+  '''Draws 3-axis pose over image
+  '''  
+  cv.line(im, pt_nose, tuple(image_pts['pitch'].ravel()), pose_types['pitch'], 3)
+  cv.line(im, pt_nose, tuple(image_pts['yaw'].ravel()), pose_types['yaw'], 3)
+  cv.line(im, pt_nose, tuple(image_pts['roll'].ravel()), pose_types['roll'], 3)
+
+
+def draw_degrees(im, pose_data, color=(0,255,0)):
+  '''Draws degrees as text over image
+  '''
+  for i, pose_type in enumerate(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, clr, thickness=2, lineType=2)
+
+
+# ---------------------------------------------------------------------------
+#
+# Matplotlib drawing functions
+#
+# ---------------------------------------------------------------------------
+
+def plot_landmarks3D(im, points, radius=3, color=(0,255,0), stroke_weight=2):
+  '''Draws facial landmarks, either 5pt or 68pt
+  '''
+  for pt in points:
+    cv.circle(im, tuple(pt), radius, color, -1, cv.LINE_AA)
+