output_utils.py

import numpy as np
import cv2
import json

FACE_BOTTOM_OUTLINE = np.arange(0, 5)
LEFT_EYEBROW = np.arange(5, 8)
RIGHT_EYEBROW = np.arange(8, 11)
LEFT_EYE_TOP = np.arange(11, 14)
LEFT_EYE_BOTTOM = np.arange(14, 17)
RIGHT_EYE_TOP = np.arange(17, 20)
RIGHT_EYE_BOTTOM = np.arange(20, 23)
NOSE = np.array([23])
MOUTH_OUTLINE = np.arange(24, 28)

FACE_OUTLINE_LIST = [FACE_BOTTOM_OUTLINE, LEFT_EYEBROW, RIGHT_EYEBROW]
LEFT_EYE_LIST = [LEFT_EYE_TOP, LEFT_EYE_BOTTOM]
RIGHT_EYE_LIST = [RIGHT_EYE_TOP, RIGHT_EYE_BOTTOM]
NOSE_LIST = [NOSE]
MOUTH_OUTLINE_LIST = [MOUTH_OUTLINE]

# (indices, BGR color, is_closed)
CONTOURS = [
    (FACE_OUTLINE_LIST, (0, 170, 255), False, "face_outline"),
    (LEFT_EYE_LIST, (50, 220, 255), False, "left_eye"),
    (RIGHT_EYE_LIST, (50, 220, 255), False, "right_eye"),
    (NOSE_LIST, (255, 30, 30), False, "nose"),
    (MOUTH_OUTLINE_LIST, (255, 30, 30), True, "mouth_outline"),
]


def visualize_box(image,
                  box,
                  score,
                  lt,
                  box_color=(0, 255, 0),
                  text_color=(255, 255, 255),
                  show_box_score=True):
    cv2.rectangle(image, tuple(box[:2]), tuple(box[2:]), box_color, lt)
    if not show_box_score:
        return
    cv2.putText(
        image,
        f'{round(score * 100, 2)}%', (box[0], box[1] - 2),
        0,
        lt / 2,
        text_color,
        thickness=max(lt, 1),
        lineType=cv2.LINE_AA)


def visualize_landmarks(image, pts, lt, landmark_score_threshold):
    for *pt, score in pts:
        pt = tuple(np.round(pt).astype(int))
        if score < landmark_score_threshold:
            color = (0, 255, 255)
        else:
            color = (0, 0, 255)
        cv2.circle(image, pt, lt, color, cv2.FILLED)


def draw_polyline(image, pts, color, closed, lt, skip_contour_with_low_score,
                  score_threshold):
    if skip_contour_with_low_score and (pts[:, 2] < score_threshold).any():
        return
    pts = np.round(pts[:, :2]).astype(int)
    cv2.polylines(image, np.array([pts], dtype=np.int32), closed, color, lt)


def visualize_contour(image, pts, lt, skip_contour_with_low_score,
                      score_threshold):
    for indices_list, color, closed, name in CONTOURS:
        for indices in indices_list:
            draw_polyline(image, pts[indices], color, closed, lt,
                          skip_contour_with_low_score, score_threshold)


def visualize(
        image,
        keypoints, bboxes,
        landmark_score_threshold,
        show_box_score=True,
        draw_contour=True,
        skip_contour_with_low_score=False):
    for pred_pts, bbox in zip(keypoints, bboxes):
        bbox, score = bbox[:4], bbox[4]
        bbox = np.round(bbox).astype(int)

        # line_thickness
        lt = max(2, int(3 * (bbox[2:] - bbox[:2]).max() / 256))

        visualize_box(image, bbox, score, lt, show_box_score=show_box_score)
        if draw_contour:
            visualize_contour(
                image,
                pred_pts,
                lt,
                skip_contour_with_low_score=skip_contour_with_low_score,
                score_threshold=landmark_score_threshold)
        visualize_landmarks(image, pred_pts, lt, landmark_score_threshold)

    return image


def save_json(
        json_file,
        keypoints,
        bboxes):
    results = []
    for pred_pts, bbox in zip(keypoints, bboxes):
        r = {}
        bbox, score = bbox[:4].tolist(), float(bbox[4])
        r['bbox'] = bbox
        r['score'] = score

        r['contours'] = {}
        for indices_list, color, closed, name in CONTOURS:
            c = []
            for indices in indices_list:
                cc = []
                for pts in pred_pts[indices]:
                    cc.append({'pos': pts[:2].tolist(), 'score': float(pts[2])})
                c.append(cc)
            r['contours'][name] = c

        results.append(r)

    with open(json_file, 'w') as f:
        json.dump(results, f, indent=2)