common/utils/visualize.py

import math
import os.path as osp
from enum import IntEnum

import numpy as np
from typing import List, Union
import cv2
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from PIL import Image, ImageDraw, ImageFont
import colorsys

from .cv import rgba_to_rgb_fixbg, img_alpha_blending
from .io_utils import imglist2imgrid, load_image


def show_factorization_on_image(img: np.ndarray,
                                explanations: np.ndarray,
                                colors: List[np.ndarray] = None,
                                image_weight: float = 0.5,
                                concept_labels: List = None,
                                visible_mask=None,) -> np.ndarray:
    """
    Modified from pytorch_grad_cam.utils.image.show_factorization_on_image

    Color code the different component heatmaps on top of the image.
        Every component color code will be magnified according to the heatmap itensity
        (by modifying the V channel in the HSV color space),
        and optionally create a lagend that shows the labels.

        Since different factorization component heatmaps can overlap in principle,
        we need a strategy to decide how to deal with the overlaps.
        This keeps the component that has a higher value in it's heatmap.

    :param img: The base image RGB format.
    :param explanations: A tensor of shape num_componetns x height x width, with the component visualizations.
    :param colors: List of R, G, B colors to be used for the components.
                   If None, will use the gist_rainbow cmap as a default.
    :param image_weight: The final result is image_weight * img + (1-image_weight) * visualization.
    :concept_labels: A list of strings for every component. If this is paseed, a legend that shows
                     the labels and their colors will be added to the image.
    :returns: The visualized image.
    """
    concept_per_pixel = explanations.argmax(axis=0)
    counts_ids = np.unique(concept_per_pixel)
    counts_ids = list(c for c in counts_ids)

    n_components = len(counts_ids)
    if colors is None:
        # taken from https://github.com/edocollins/DFF/blob/master/utils.py
        _cmap = plt.cm.get_cmap('gist_rainbow')
        colors = [
            np.array(
                _cmap(i)) for i in np.arange(
                0,
                1,
                1.0 /
                n_components)]

    masks = []
    for ii, counts_id in enumerate(counts_ids):
        mask = np.zeros(shape=(img.shape[0], img.shape[1], 3))
        mask[:, :, :] = colors[ii][:3]
        explanation = explanations[counts_id]
        # explanation[concept_per_pixel != i] = 0
        explanation = concept_per_pixel == counts_id
        mask = np.uint8(mask * 255)
        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2HSV)
        mask[:, :, 2] = np.uint8(255 * explanation)
        mask = cv2.cvtColor(mask, cv2.COLOR_HSV2RGB)
        mask = np.float32(mask) / 255
        masks.append(mask)

    mask = np.sum(np.float32(masks), axis=0)
    result = img * image_weight + mask * (1 - image_weight)
    result = np.uint8(result * 255)
    if visible_mask is not None:
        result = result * visible_mask + np.full_like(result, fill_value=255) * (1 - visible_mask)
        result = result.astype(np.uint8)

    if concept_labels is not None:
        px = 1 / plt.rcParams['figure.dpi']  # pixel in inches
        fig = plt.figure(figsize=(result.shape[1] * px, result.shape[0] * px))
        plt.rcParams['legend.fontsize'] = int(
            20 * result.shape[0] / 256 / max(1, n_components / 6))
        lw = 5 * result.shape[0] / 256
        lines = [Line2D([0], [0], color=colors[i], lw=lw)
                 for i in range(n_components)]
        c_labels = [concept_labels[i] for i in counts_ids]
        plt.legend(lines,
                   c_labels,
                   mode="expand",
                   fancybox=True,
                   shadow=True)

        plt.tight_layout(pad=0, w_pad=0, h_pad=0)
        plt.axis('off')
        fig.canvas.draw()
        data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
        plt.close(fig=fig)
        data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
        data = cv2.resize(data, (result.shape[1], result.shape[0]))
        result = np.hstack((result, data))
    return result


def visualize_segs_with_labels(mask_list, src_img: np.ndarray, tag_list, image_weight=0.3, colors=None, reference_img: np.ndarray = None, draw_legend=True) -> np.ndarray:

    n_components = len(mask_list)

    img = rgba_to_rgb_fixbg(src_img)

    colors = np.array([get_color(ii) for ii in range(n_components)], dtype=np.uint8)
    
    colored_list = [img]
    for ii, mask in enumerate(mask_list):
        color = np.array(colors[ii][:3])
        mask = mask.astype(np.float32)
        c = np.zeros((img.shape[0], img.shape[1], 4), dtype=np.uint8)
        c[..., :3] = (color[None, None] * mask[..., None]).astype(np.uint8)
        c[..., -1] = np.clip(mask.astype(np.float32) * 255 * (1 - image_weight), 0, 255).astype(np.uint8)
        colored_list.append(c)

    colored_final = img_alpha_blending(colored_list)
    result = colored_final[..., :3]

    if not draw_legend or tag_list is None:
        if reference_img is not None:
            result = np.hstack((reference_img, result))
        return result

    colors = colors.astype(np.float32) / 255.
    px = 1 / plt.rcParams['figure.dpi']  # pixel in inches
    fig = plt.figure(figsize=(result.shape[1] * px, result.shape[0] * px), facecolor=[0, 0, 0, 0])
    
    fnt_sz = int(5 * result.shape[0] / 256)
    plt.rcParams['legend.fontsize'] = fnt_sz
    lw = 5 * result.shape[0] / 256
    lines = [Line2D([0], [0], color=colors[i], lw=lw)
                for i in range(n_components)]
    # c_labels = [all_labels[i] for i in all_labels]
    plt.legend(lines,
                tag_list,
                mode="expand",
                fancybox=False,
                edgecolor="black",
                # frameon=False,
                shadow=False,
                framealpha=0.)

    plt.tight_layout(pad=0, w_pad=0, h_pad=0)
    plt.axis('off')
    fig.canvas.draw()
    data = np.frombuffer(fig.canvas.buffer_rgba() , dtype=np.uint8)
    plt.close(fig=fig)
    data = data.reshape(fig.canvas.get_width_height()[::-1] + (4,))
    dx, dy, dw, dh = cv2.boundingRect(cv2.findNonZero(data[..., 3]))
    
    data = rgba_to_rgb_fixbg(data[:, dx: dx + dw])
    data = cv2.copyMakeBorder(data, 0, 0, fnt_sz, fnt_sz, borderType=cv2.BORDER_CONSTANT, value=(255, 255, 255))
    if data.shape[0] != img.shape[0]:
        data = cv2.resize(data, (dw, result.shape[0]))

    if reference_img is not None:
        result = np.hstack((reference_img, result, data))
    else:
        result = np.hstack((result, data))
    return result


def pixel_rounder(n, mode):
    if mode==True or mode=='round':
        return round(n)
    elif mode=='ceil':
        return math.ceil(n)
    elif mode=='floor':
        return math.floor(n)
    else:
        return n


def pixel_ij(x, rounding=True):
    if isinstance(x, np.ndarray):
        x = x.tolist()
    return tuple(pixel_rounder(i, rounding) for i in (
        x if isinstance(x, tuple) or isinstance(x, list) else (x,x)
    ))


def ucolors(num_colors):
    # uniform color generator
    colors=[]
    for i in np.arange(0., 360., 360. / num_colors):
        hue = i/360.
        lightness = (50 + np.random.rand() * 10)/100.
        saturation = (90 + np.random.rand() * 10)/100.
        colors.append(colorsys.hls_to_rgb(hue, lightness, saturation))
    return colors

coco_parts_dict = {
    "ear_right-eye_right": (4, 2),
    "eye_right-nose": (2, 0),
    "nose-eye_left": (0, 1),
    "eye_left-ear_left": (1, 3),

    "wrist_right-elbow_right": (10, 8),
    "elbow_right-shoulder_right": (8, 6),
    "shoulder_right-shoulder_left": (6, 5),
    "shoulder_left-elbow_left": (5, 7),
    "elbow_left-wrist_left": (7, 9),

    "ankle_right-knee_right": (16, 14),
    "knee_right-hip_right": (14, 12),
    "hip_right-hip_left": (12, 11),
    "hip_left-knee_left": (11, 13),
    "knee_left-ankle_left": (13, 15)
}

coco_parts = [
    (4, 2),  #  0, ear_right - eye_right
    ( 2, 0),  #  1, eye_right - nose
    ( 0, 1),  #  2, nose - eye_left
    ( 1, 3),  #  3, eye_left - ear_left

    (10, 8),  #  4, wrist_right - elbow_right
    ( 8, 6),  #  5, elbow_right - shoulder_right
    ( 6, 5),  #  6, shoulder_right - shoulder_left
    ( 5, 7),  #  7, shoulder_left - elbow_left
    ( 7, 9),  #  8, elbow_left - wrist_left

    (16,14),  #  9, ankle_right - knee_right
    (14,12),  # 10, knee_right - hip_right
    (12,11),  # 11, hip_right - hip_left
    (11,13),  # 12, hip_left - knee_left
    (13,15),  # 13, knee_left- ankle_left
]
coco_parts_ext = coco_parts + [
    ( 0,17),  # 14, nose - nose_root
    ( 9,19),  # 15, wrist_left - thumb_left
    (10,20),  # 16, wrist_right - thumb_right
    (15,21),  # 17, ankle_left - toe_left
    (16,22),  # 18, ankle_right - toe_right
]

coco_part_colors = ucolors(len(coco_parts))
coco_part_colors_ext = coco_part_colors + ucolors(len(coco_parts_ext)-len(coco_parts))

coco_keypoints = [
    'nose',             #  0
    'eye_left',         #  1
    'eye_right',        #  2
    'ear_left',         #  3
    'ear_right',        #  4

    'shoulder_left',    #  5
    'shoulder_right',   #  6
    'elbow_left',       #  7
    'elbow_right',      #  8
    'wrist_left',       #  9
    'wrist_right',      # 10

    'hip_left',         # 11
    'hip_right',        # 12
    'knee_left',        # 13
    'knee_right',       # 14
    'ankle_left',       # 15
    'ankle_right',      # 16
]

coco_keypoints_ext = coco_keypoints + [
    'nose_root',    # 17
    'body_upper',   # 18
    'thumb_left',   # 19
    'thumb_right',  # 20
    'toe_left',     # 21
    'toe_right',    # 22
]

def c255(c):
    # color format utility
    if c is None:
        return None
    if isinstance(c, str):
        c = {
            'r': (1,0,0),
            'g': (0,1,0),
            'b': (0,0,1),
            'k': 0,
            'w': 1,
            't': (0,1,1),
            'm': (1,0,1),
            'y': (1,1,0),
            'a': (0,0,0,0),
        }[c]
    if isinstance(c, list) or isinstance(c, tuple):
        if len(c)==3:
            c = c + (1,)
        elif len(c)==1:
            c = (c,c,c,1)
        c = tuple(int(255*q) for q in c)
    else:
        c = int(255*c)
        c = (c,c,c,255)
    return c

def draw_rect(image, corner, size, w=1, c='r', f=None):
    corner = pixel_ij(corner, rounding=True)
    size = pixel_ij(size, rounding=True)
    w = max(1, round(w))
    c = c255(c)
    f = c255(f)
    ans = image.copy()
    d = ImageDraw.Draw(ans)
    d.rectangle(
        [corner[1], corner[0], corner[1]+size[1]-1, corner[0]+size[0]-1],
        fill=f, outline=c, width=w,
    )
    return ans

def draw_dot(d: ImageDraw.ImageDraw, point, s=1, c='r'):
    c = c255(c)
    x,y = pixel_ij(point, rounding=False)
    d.ellipse(
        [(y-s,x-s), (y+s,x+s)],
        fill=c,
    )

def draw_line(d: ImageDraw.ImageDraw, a, b, w=1, c='r'):
    a = pixel_ij(a, rounding=False)
    b = pixel_ij(b, rounding=False)
    c = c255(c)
    w = max(1, round(w))
    d.line([a[::-1], (b[1]-1,b[0]-1)], fill=c, width=w)

    

def visualize_pos_keypoints(image: Image.Image, bbox=None, keypoints=None, scores=None):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    if bbox is not None:
        image = draw_rect(image, *bbox, c='r', w=2)
    image_rst = image.copy()
    d = ImageDraw.Draw(image_rst)
    if keypoints is not None:
        if isinstance(keypoints, dict):
            keypoints = np.asarray([keypoints[k] for k in coco_keypoints])
        for (a,b),c in zip(coco_parts, coco_part_colors):
            draw_line(d, keypoints[a], keypoints[b], w=5, c=c)
        keypoints = keypoints[:len(coco_keypoints)]
        for ii, kp in enumerate(keypoints):
            draw_dot(d, kp, s=5, c='r')
        font = ImageFont.truetype("assets/arial.ttf", size=20)
        if scores is not None:
            # fnt = ImageFont.load_default(size=30)
            for ii, kp in enumerate(keypoints):
                s = scores[ii]
                d.text((kp[1], kp[0]), f'{int(round(s * 100))}', font=font, fill=(0, 0, 0, 255), stroke_fill=(255, 255, 255, 255), stroke_width=2)

    return image_rst



def pil_draw_text(image: Image.Image, text, point, fill=(0, 0, 0, 255), font_size=20, stroke_fill=(255, 255, 255, 255), font=None, stroke_width=0):
    is_ndarray = False
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
        is_ndarray = True
    d = ImageDraw.Draw(image)
    if font is None:
        font = ImageFont.truetype("assets/arial.ttf", size=font_size)
    if not isinstance(text, list):
        text = [text]
        point = [point]
    for t, p in zip(text, point):
        d.text(p, t, font=font, fill=fill, stroke_fill=stroke_fill, stroke_width=stroke_width)
    if is_ndarray:
        image = np.array(image)
    return image




# https://github.com/hysts/anime-face-detector/blob/main/assets/landmarks.jpg
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)
FACE_CONTOURS = [
    (FACE_OUTLINE_LIST, (0, 170, 255), False),
    (LEFT_EYE_LIST, (50, 220, 255), False),
    (RIGHT_EYE_LIST, (50, 220, 255), False),
    (NOSE_LIST, (255, 30, 30), False),
    (MOUTH_OUTLINE_LIST, (255, 30, 30), True),
]

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_face_contour(image, pts, lt, skip_contour_with_low_score,
                      score_threshold):
    for indices_list, color, closed in FACE_CONTOURS:
        for indices in indices_list:
            draw_polyline(image, pts[indices], color, closed, lt,
                          skip_contour_with_low_score, score_threshold)



def visualize_facedet_output(image: np.ndarray,
              preds: np.ndarray,
              face_score_threshold: float = 0.5,
              landmark_score_threshold: float = 0.3,
              show_box_score: bool = True,
              draw_contour: bool = True,
              skip_contour_with_low_score=False):
    res = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    for pred in preds:
        box = pred['bbox']
        box, score = box[:4], box[4]
        box = np.round(box).astype(int)
        pred_pts = pred['keypoints']

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

        visualize_box(res, box, score, lt, show_box_score=show_box_score)
        if draw_contour:
            visualize_face_contour(
                res,
                pred_pts,
                lt,
                skip_contour_with_low_score=skip_contour_with_low_score,
                score_threshold=landmark_score_threshold)
        visualize_landmarks(res, pred_pts, lt, landmark_score_threshold)

    res = cv2.cvtColor(res, cv2.COLOR_BGR2RGB)
    return res



FACE_LABEL2NAME = {
    0:'background', 
    1:'skin', # it's face
    2:'l_brow', 3:'r_brow',
    4:'l_eye', 5:'r_eye', 
    6:'eye_g',  # glass 
    7:'l_ear', 8:'r_ear', 
    9:'ear_r',  # not used
    10:'nose', 11:'mouth',
    12:'u_lip', 13:'l_lip', # both not used 
    14:'neck', 
    15:'neck_l',    # not used
    16:'cloth',     # mostly collar
    17:'hair', 
    18:'hat'
}

VALID_FACE_GROUPS = {
    'face': [1],
    'brows': [2, 3],
    'eyes': [4, 5],
    'glass': [6],
    'ears': [7, 8],
    'nose': [10],
    'mouth': [11],
    'neck': [14],
    'cloth': [16],
    'hair': [17],
    'hat': [18]
}


def visualize_segs(labels: Union[np.ndarray, List], src_img: np.ndarray = None, image_weight = 0.5, colors=None, label2name_map=None, output_dtype='numpy'):
    '''
    labels: (H, W) components map or (N, H, W) masks
    '''

    if label2name_map is None:
        label2name_map = {l: str(l) for l in range(len(labels))}

    if isinstance(labels, list):
        labels = [label[None] if label.ndim == 2 else label for label in labels]
        labels = np.concatenate(labels)

    # if colors is None:
    #     colors = np.array([get_color(ii) for ii in range(len(label2name_map))], dtype=np.float32) / 255.
    
    # h, w = labels.shape[-2:]
    # vis = np.zeros((h, w, 3), dtype=np.float32)

    # if labels.ndim == 3:
    #     assert labels.shape[0] == len(label2name_map)
    # else:
    #     assert labels.ndim == 2

    # for label in label2name_map:
    #     if labels.ndim == 2:
    #         mask = (labels == label).astype(np.float32)[..., None]
    #     else:
    #         mask = labels[label].astype(np.float32)[..., None]
    #     vis += mask * colors[label][:3].reshape((1, 1, -1))

    # if src_img is not None:
    #     if np.max(src_img) > 1:
    #         src_img = src_img.astype(np.float32) / 255.
    #     vis = image_weight * src_img + (1 - image_weight) * vis

    n_components = len(labels)
    colors = np.array([get_color(ii) for ii in range(n_components)], dtype=np.uint8)

    img = src_img
    mask_list = labels
    colored_list = [src_img]
    mask_full = np.zeros_like(mask_list[0]).astype(np.float32)
    c = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
    for ii, mask in enumerate(mask_list):
        color = np.array(colors[ii][:3])
        mask_f = mask.astype(np.float32)
        mask_full += mask_f
        # c = np.zeros((img.shape[0], img.shape[1], 4), dtype=np.uint8)
        c[np.where(mask)] = (color).astype(np.uint8)
        # c[..., -1] = np.clip(mask.astype(np.float32) * 255 * (1 - image_weight), 0, 255).astype(np.uint8)
        # colored_list.append(c)

    # colored_final = img_alpha_blending(colored_list)
    mask_full = np.clip(mask_full[..., None], 0, 1) * (1-image_weight)
    vis = np.round(c * mask_full + img * (1-mask_full)).astype(np.uint8)

    # vis = (np.clip(vis, 0, 1) * 255).astype(np.uint8)
    if output_dtype.lower() == 'pil':
        return Image.fromarray(vis)
    return vis


COLOR_PALETTE = [
    (220, 20, 60), (119, 11, 32), (0, 0, 142), (0, 0, 230), (106, 0, 228),
    (0, 60, 100), (0, 80, 100), (0, 0, 70), (0, 0, 192), (250, 170, 30),
    (100, 170, 30), (220, 220, 0), (175, 116, 175), (250, 0, 30),
    (165, 42, 42), (255, 77, 255), (0, 226, 252), (182, 182, 255),
    (0, 82, 0), (120, 166, 157), (110, 76, 0), (174, 57, 255),
    (199, 100, 0), (72, 0, 118), (255, 179, 240), (0, 125, 92),
    (209, 0, 151), (188, 208, 182), (0, 220, 176), (255, 99, 164),
    (92, 0, 73), (133, 129, 255), (78, 180, 255), (0, 228, 0),
    (174, 255, 243), (45, 89, 255), (134, 134, 103), (145, 148, 174),
    (255, 208, 186), (197, 226, 255), (171, 134, 1), (109, 63, 54),
    (207, 138, 255), (151, 0, 95), (9, 80, 61), (84, 105, 51),
    (74, 65, 105), (166, 196, 102), (208, 195, 210), (255, 109, 65),
    (0, 143, 149), (179, 0, 194), (209, 99, 106), (5, 121, 0),
    (227, 255, 205), (147, 186, 208), (153, 69, 1), (3, 95, 161),
    (163, 255, 0), (119, 0, 170), (0, 182, 199), (0, 165, 120),
    (183, 130, 88), (95, 32, 0), (130, 114, 135), (110, 129, 133),
    (166, 74, 118), (219, 142, 185), (79, 210, 114), (178, 90, 62),
    (65, 70, 15), (127, 167, 115), (59, 105, 106), (142, 108, 45),
    (196, 172, 0), (95, 54, 80), (128, 76, 255), (201, 57, 1),
    (246, 0, 122), (191, 162, 208), (255, 255, 128), (147, 211, 203),
    (150, 100, 100), (168, 171, 172), (146, 112, 198), (210, 170, 100),
    (92, 136, 89), (218, 88, 184), (241, 129, 0), (217, 17, 255),
    (124, 74, 181), (70, 70, 70), (255, 228, 255), (154, 208, 0),
    (193, 0, 92), (76, 91, 113), (255, 180, 195), (106, 154, 176),
    (230, 150, 140), (60, 143, 255), (128, 64, 128), (92, 82, 55),
    (254, 212, 124), (73, 77, 174), (255, 160, 98), (255, 255, 255),
    (104, 84, 109), (169, 164, 131), (225, 199, 255), (137, 54, 74),
    (135, 158, 223), (7, 246, 231), (107, 255, 200), (58, 41, 149),
    (183, 121, 142), (255, 73, 97), (107, 142, 35), (190, 153, 153),
    (146, 139, 141), (70, 130, 180), (134, 199, 156), (209, 226, 140),
    (96, 36, 108), (96, 96, 96), (64, 170, 64), (152, 251, 152),
    (208, 229, 228), (206, 186, 171), (152, 161, 64), (116, 112, 0),
    (0, 114, 143), (102, 102, 156), (250, 141, 255)
]

class Colors:
    # Ultralytics color palette https://ultralytics.com/
    def __init__(self):
        # hex = matplotlib.colors.TABLEAU_COLORS.values()
        # hexs = ('FF1010', '10FF10', 'FFF010', '100FFF', 'c0c0c0', 'FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
        #         '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        hexs = [
            '#4363d8',
            '#9A6324',
            '#808000',
            '#469990',
            '#000075',
            '#e6194B',
            '#f58231',
            '#ffe119',
            '#bfef45',
            '#3cb44b',
            '#42d4f4',
            '#800000',
            '#911eb4',
            '#f032e6',
            '#fabed4',
            '#ffd8b1',
            '#fffac8',
            '#aaffc3',
            '#dcbeff',
            '#a9a9a9',
            '#006400',
            '#4169E1',
            '#8B4513',
            '#FA8072',
            '#87CEEB',
            '#FFD700',
            '#ffffff',
            '#000000',
        ]
        self.palette = [self.hex2rgb(f'#{c}') if not c.startswith('#') else self.hex2rgb(c) for c in hexs]
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  # rgb order (PIL)
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
    
DEFAULT_COLOR_PALETTE = Colors()
def get_color(idx):
    if idx == -1:
        return 255
    else:
        return DEFAULT_COLOR_PALETTE(idx)



def vis_parts(srcd: str, tag_list, nmax_samples=12, cols=4):
    partsd = osp.join(srcd, 'parts')
    rst_list = []
    nparts = 0
    for tag in tag_list:
        p = osp.join(partsd, tag + '_vis.png')
        if not osp.exists(p):
            continue
        img = Image.open(p)
        pil_draw_text(img, tag, point=(0, 0), font_size=128, stroke_width=12)
        rst_list.append(img)
        if len(rst_list) >= nmax_samples:
            vis = imglist2imgrid(rst_list, cols=cols)
            Image.fromarray(vis).save(osp.join(srcd, f'part_vis{nparts}.jpg'), q=97)
            rst_list = []
            nparts += 1

    if len(rst_list) > 0:
        vis = imglist2imgrid(rst_list, cols=cols)
        Image.fromarray(vis).save(osp.join(srcd, f'part_vis{nparts}.jpg'), q=97)


def imglist2imgrid_with_tags(
        img_list, tag_list, cols=4, output_type='numpy', skip_empty=False, fix_size=None, tag_breakn=4, font_size=-1, stroke_width=-1, font_pos = (0, 0)):
    
    def _wrap_text(tags):
        '''
        maybe wrap according to actual render size
        '''
        if isinstance(tags, str):
            tags = tags.split(',')
        caption = ''
        nadded = 0
        for t in tags:
            if nadded >= tag_breakn:
                nadded = 1
                caption = caption + '\n' + t
            else:
                nadded += 1
                if caption == '':
                    caption = t
                else:
                    caption = caption + ',' + t
        return caption

    new_img_lst = []
    for ii, img in enumerate(img_list):
        if skip_empty and tag_list[ii].strip() == '':
            continue

        if isinstance(img, np.ndarray):
            img = Image.fromarray(img)
        else:
            assert isinstance(img, Image.Image)
            img = img.copy()
        ih, iw = img.height, img.width
        if font_size < 0:
            font_size = max(ih, iw) // 16
        if stroke_width < 0:
            stroke_width = font_size // 10

        pil_draw_text(img, _wrap_text(tag_list[ii]), point=font_pos, font_size=font_size, stroke_width=stroke_width)
        new_img_lst.append(img)
        
    return imglist2imgrid(new_img_lst, cols=cols, output_type=output_type, fix_size=fix_size)



class JointType(IntEnum):
    Nose = 0
    Neck = 1
    RightShoulder = 2
    RightElbow = 3
    RightHand = 4
    LeftShoulder = 5
    LeftElbow = 6
    LeftHand = 7
    RightWaist = 8
    RightKnee = 9
    RightFoot = 10
    LeftWaist = 11
    LeftKnee = 12
    LeftFoot = 13
    RightEye = 14
    LeftEye = 15
    RightEar = 16
    LeftEar = 17
    RightToes = 18
    LeftToes = 19
    RightFist = 20
    LeftFist = 21
    Spine = 22


def plot_faceparsing(img: np.ndarray, pose, face_landmarks):

    landmark_params = {
        'moe_joint_indices': {
            'nose': JointType.Nose,
            'L_eye': JointType.LeftEye,
            'R_eye': JointType.RightEye,
            'L_ear': JointType.LeftEar,
            'R_ear': JointType.RightEar,
            'neck': JointType.Neck,
            'L_shoulder': JointType.LeftShoulder,
            'R_shoulder': JointType.RightShoulder,
            'L_elbow': JointType.LeftElbow,
            'R_elbow': JointType.RightElbow,
            'L_hand': JointType.LeftHand,
            'R_hand': JointType.RightHand,
            'L_waist': JointType.LeftWaist,
            'R_waist': JointType.RightWaist,
            'L_knee': JointType.LeftKnee,
            'R_knee': JointType.RightKnee,
            'L_foot': JointType.LeftFoot,
            'R_foot': JointType.RightFoot
        },
        'limbs_point_plot': [
            ['neck', 'R_waist'],
            ['R_waist', 'R_knee'],
            ['R_knee', 'R_foot'],
            ['neck', 'L_waist'],
            ['L_waist', 'L_knee'],
            ['L_knee', 'L_foot'],
            ['neck', 'R_shoulder'],
            ['R_shoulder', 'R_elbow'],
            ['R_elbow', 'R_hand'],
            ['R_shoulder', 'R_ear'],
            ['neck', 'L_shoulder'],
            ['L_shoulder', 'L_elbow'],
            ['L_elbow', 'L_hand'],
            ['L_shoulder', 'L_ear'],
            ['neck', 'nose'],
            ['nose', 'R_eye'],
            ['nose', 'L_eye'],
            ['R_eye', 'R_ear'],
            ['L_eye', 'L_ear']
        ]
    }

    landmark_all_groups = ['outline',
                        'nose',
                        'mouth',
                        'left_eye_brow',
                        'right_eye_brow',
                        'left_eye_outline',
                        'right_eye_outline',
                        'right_eye',
                        'left_eye']

    def _plot_pose(orig_img, pose, line_width=2, circle_size=3):
        limb_colors = [
            [0, 255, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255],
            [0, 85, 255], [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0.],
            [255, 0, 85], [170, 255, 0], [85, 255, 0], [170, 0, 255.], [0, 0, 255],
            [0, 0, 255], [255, 0, 255], [170, 0, 255], [255, 0, 170],
        ]
        joint_colors = [
            [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0],
            [85, 255, 0], [0, 255, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255],
            [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], [170, 0, 255],
            [255, 0, 255], [255, 0, 170], [255, 0, 85]]

        canvas = orig_img

        for i, (limb, color) in enumerate(zip(landmark_params['limbs_point_plot'], limb_colors)):
            if i != 9 and i != 13:  # don't show ear-shoulder connection
                keya = limb[0]
                keyb = limb[1]
                xa = pose['pose'][keya]['x']
                ya = pose['pose'][keya]['y']
                va = pose['pose'][keya]['v']
                xb = pose['pose'][keyb]['x']
                yb = pose['pose'][keyb]['y']
                vb = pose['pose'][keyb]['v']
                if va > 0 and vb > 0:
                    cv2.line(canvas, (xa, ya), (xb, yb), color, line_width)

        for key, color in zip(list(landmark_params['moe_joint_indices'].keys()), joint_colors):
            x = pose['pose'][key]['x']
            y = pose['pose'][key]['y']
            v = pose['pose'][key]['v']
            if v != 0:
                cv2.circle(canvas, (x, y), circle_size, color, -1)
        return canvas

    def _plot_face(orig_img, landmarks, draw_line=True, line_width=2, circle_size=3):
        
        group_line_color = {
            'outline': [0, 255, 0],
            'left_eye_brow': [0, 255, 85], 
            'right_eye_brow': [0, 255, 170], 
            'left_eye_outline': [0, 255, 255], 
            'right_eye_outline': [0, 170, 255],
            'nose': [0, 85, 255],
            'mouth': [255, 85, 0]
        }
        
        group_colors = {
            'nose': [0, 0, 0],
            'mouth': [127, 127, 127],
            'left_eye': [0, 127, 0],
            'right_eye': [0, 0, 127],
            'outline': [255, 0, 0],
            'left_eye_outline': [255, 0, 255],
            'right_eye_outline': [255, 255, 0],
            'left_eye_brow': [0, 255, 0],
            'right_eye_brow': [0, 0, 255],
        }

        group_line = {
            'outline': [[0, 1],
                        [1, 2],
                        [2, 3],
                        [3, 4],
                        [4, 5],
                        [5, 6],
                        [6, 7],
                        [7, 8]],
            'left_eye_brow': [[0, 1], [1, 2]],
            'right_eye_brow': [[0, 1], [1, 2]],
            'left_eye_outline': [[0, 1],
                                [1, 2],
                                [2, 3],
                                [3, 4],
                                [4, 5],
                                [5, 0]],
            'right_eye_outline': [[0, 1],
                                [1, 2],
                                [2, 3],
                                [3, 4],
                                [4, 5],
                                [5, 0]],
            'nose': [[0, 1]],
            'mouth': [[0, 1],
                    [1, 2],
                    [2, 3],
                    [3, 4],
                    [4, 5],
                    [5, 0]],
        }

        canvas = orig_img
        if draw_line:
            for k, color in group_line_color.items():
                if k in landmarks:
                    for limb in group_line[k]:
                        keya = limb[0]
                        keyb = limb[1]

                        xa = int(landmarks[k][keya][0])
                        ya = int(landmarks[k][keya][1])
                        xb = int(landmarks[k][keyb][0])
                        yb = int(landmarks[k][keyb][1])

                        cv2.line(canvas, (xa, ya), (xb, yb), color, line_width)

        for k, color in group_colors.items():
            if k in landmarks:
                for points in landmarks[k]:
                    x = int(points[0])
                    y = int(points[1])
                    cv2.circle(canvas, (x, y), circle_size, color, -1)
        return canvas

    def shift_position(pos, left, top):
        new_pos = pos.copy()
        new_pos[:, 0] += left
        new_pos[:, 1] += top
        return new_pos

    canvas = img.copy()
    _plot_pose(canvas, pose)

    def get_points(group_names=[]):
        if isinstance(group_names, list):
            ret = {}
            for g in group_names:
                ret[g] = shift_position(face_landmarks['points'][g], face_landmarks['left'], face_landmarks['top'])
            return ret

        return shift_position(face_landmarks['points'][group_names], face_landmarks['left'], face_landmarks['top'])

    pos = np.array(face_landmarks['points'])
    
    _plot_face(canvas, get_points(landmark_all_groups))
    return canvas


RECT_LINES = ((0, 1), (1, 2), (2, 3), (3, 0))

def uint82bin(n, count=8):
    """returns the binary of integer n, count refers to amount of bits"""
    return ''.join([str((n >> y) & 1) for y in range(count - 1, -1, -1)])


def labelcolormap(N):
    if N == 35:  # cityscape
        cmap = np.array([(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (111, 74, 0), (81, 0, 81),
                         (128, 64, 128), (244, 35, 232), (250, 170, 160), (230, 150, 140), (70, 70, 70), (102, 102, 156), (190, 153, 153),
                         (180, 165, 180), (150, 100, 100), (150, 120, 90), (153, 153, 153), (153, 153, 153), (250, 170, 30), (220, 220, 0),
                         (107, 142, 35), (152, 251, 152), (70, 130, 180), (220, 20, 60), (255, 0, 0), (0, 0, 142), (0, 0, 70),
                         (0, 60, 100), (0, 0, 90), (0, 0, 110), (0, 80, 100), (0, 0, 230), (119, 11, 32), (0, 0, 142)],
                        dtype=np.uint8)
    else:
        cmap = np.zeros((N, 3), dtype=np.uint8)
        for i in range(N):
            r, g, b = 0, 0, 0
            id = i + 1  # let's give 0 a color
            for j in range(7):
                str_id = uint82bin(id)
                r = r ^ (np.uint8(str_id[-1]) << (7 - j))
                g = g ^ (np.uint8(str_id[-2]) << (7 - j))
                b = b ^ (np.uint8(str_id[-3]) << (7 - j))
                id = id >> 3
            cmap[i, 0] = r
            cmap[i, 1] = g
            cmap[i, 2] = b

    return cmap
    

def plot_points_lines(orig_img, points, lines=RECT_LINES, line_width=2, circle_size=3, ref_pos=(0, 0)):
    colors_p = labelcolormap(len(points))
    colors_l = labelcolormap(len(lines))

    c = orig_img.shape[-1]
    
    canvas = orig_img.copy()
    for i in range(len(lines)):
        keya = lines[i][0]
        keyb = lines[i][1]
        xa = int(points[keya][0] + ref_pos[0])
        ya = int(points[keya][1] + ref_pos[1])
        xb = int(points[keyb][0] + ref_pos[0])
        yb = int(points[keyb][1] + ref_pos[1])
        # print(colors_l[i])
        color = [int(colors_l[i][0]), int(colors_l[i][1]), int(colors_l[i][2])]
        cv2.line(canvas, (xa, ya), (xb, yb), color, line_width)
        
    for i in range(len(points)):
        x = int(points[i][0] + ref_pos[0])
        y = int(points[i][1] + ref_pos[1])
        color = [int(colors_p[i][0]), int(colors_p[i][1]), int(colors_p[i][2])]
        if c == 4:
            color.append(255)
        cv2.circle(canvas, (x, y), circle_size, color, -1)
    return canvas