diff options
Diffstat (limited to 'cli/app/utils/draw_utils.py')
| -rw-r--r-- | cli/app/utils/draw_utils.py | 215 |
1 files changed, 215 insertions, 0 deletions
diff --git a/cli/app/utils/draw_utils.py b/cli/app/utils/draw_utils.py new file mode 100644 index 0000000..18c03f2 --- /dev/null +++ b/cli/app/utils/draw_utils.py @@ -0,0 +1,215 @@ +import sys +from math import sqrt + +import numpy as np +import cv2 as cv +import PIL +from PIL import ImageDraw + +from app.utils import im_utils +from app.settings import app_cfg + +log = app_cfg.LOG + +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 +# +# --------------------------------------------------------------------------- + +pose_types = {'pitch': (0,0,255), 'roll': (255,0,0), 'yaw': (0,255,0)} + +def draw_landmarks2d_cv(im, points_norm, radius=3, color=(0,255,0)): + '''Draws facial landmarks, either 5pt or 68pt + ''' + im_dst = im.copy() + dim = im.shape[:2][::-1] + for x,y in points_norm: + pt = (int(x*dim[0]), int(y*dim[1])) + cv.circle(im_dst, pt, radius, color, -1, cv.LINE_AA) + return im_dst + +def draw_landmarks2d_pil(im_pil, points_norm, radius=3, color=(0,255,0)): + '''Draws facial landmarks, either 5pt or 68pt + ''' + assert im_utils.is_pil(im_pil) + draw = ImageDraw.Draw(im_pil) + dim = im.shape[:2][::-1] + for x,y in points_norm: + x1, y1 = (int(x*dim[0]), int(y*dim[1])) + xyxy = (x1, y1, x1+radius, y1+radius) + draw.ellipse(xyxy, fill='white') + del draw + im_dst = im_utils.ensure_np(im_pil) + im_dst = im_utils.rgb2bgr(im_dst) + return im_dst + + +def draw_landmarks3D_cv(im, points, radius=3, color=(0,255,0)): + '''Draws 3D facial landmarks + ''' + im_dst = im.copy() + for x,y,z in points: + cv.circle(im_dst, (x,y), radius, color, -1, cv.LINE_AA) + return im_dst + + +def draw_bbox_cv(im_np, bbox_norm, color=(0,255,0), stroke_weight=2): + '''Draws BBox onto cv image + ''' + bbox_dim = bbox_norm.to_bbox_dim(im_np.shape[:2][::-1]) + return cv.rectangle(im_np, bbox_dim.p1.xy, bbox_dim.p2.xy, color, stroke_weight, cv.LINE_AA) + + +def draw_bbox_pil(im, bboxes_norm, color=(0,255,0), stroke_weight=2): + '''Draws BBox onto cv image + :param color: RGB value + ''' + if im_utils.is_np(im): + im = im_utils.np2pil(im) + was_np = True + else: + was_np = False + + if not type(bboxes_norm) == list: + bboxes_norm = [bboxes_norm] + + + im_draw = ImageDraw.ImageDraw(im) + + for bbox_norm in bboxes_norm: + bbox_dim = bbox_norm.to_bbox_dim(im.size) + xyxy = (bbox_dim.p1.xy, bbox_dim.p2.xy) + im_draw.rectangle(xyxy, outline=color, width=stroke_weight) + del im_draw + + if was_np: + im = im_utils.pil2np(im) + return im + + +def draw_pose(im, pt_nose, image_pts): + '''Draws 3-axis pose over image + TODO: normalize point data + ''' + im_dst = im.copy() + log.debug(f'pt_nose: {pt_nose}') + log.debug(f'image_pts pitch: {image_pts["pitch"]}') + cv.line(im_dst, pt_nose, tuple(image_pts['pitch']), pose_types['pitch'], 3) + cv.line(im_dst, pt_nose, tuple(image_pts['yaw']), pose_types['yaw'], 3) + cv.line(im_dst, pt_nose, tuple(image_pts['roll']), pose_types['roll'], 3) + return im_dst + +def draw_text_cv(im, pt_norm, text, size=1.0, color=(0,255,0)): + '''Draws degrees as text over image + ''' + im_dst = im.copy() + dim = im.shape[:2][::-1] + pt = tuple(map(int, (pt_norm[0]*dim[0], pt_norm[1]*dim[1]))) + cv.putText(im_dst, text, pt, cv.FONT_HERSHEY_SIMPLEX, size, color, thickness=1, lineType=cv.LINE_AA) + return im_dst + + +def draw_degrees(im, pose_data, color=(0,255,0)): + '''Draws degrees as text over image + ''' + im_dst = im.copy() + 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_dst, t, origin, cv.FONT_HERSHEY_SIMPLEX, 0.5, clr, thickness=2, lineType=2) + return im_dst
\ No newline at end of file |