unianimate/dwpose/util.py

import math
import numpy as np
import colorsys
import cv2


eps = 0.01


def smart_resize(x, s):
    Ht, Wt = s
    if x.ndim == 2:
        Ho, Wo = x.shape
        Co = 1
    else:
        Ho, Wo, Co = x.shape
    if Co == 3 or Co == 1:
        k = float(Ht + Wt) / float(Ho + Wo)
        return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
    else:
        return np.stack([smart_resize(x[:, :, i], s) for i in range(Co)], axis=2)


def smart_resize_k(x, fx, fy):
    if x.ndim == 2:
        Ho, Wo = x.shape
        Co = 1
    else:
        Ho, Wo, Co = x.shape
    Ht, Wt = Ho * fy, Wo * fx
    if Co == 3 or Co == 1:
        k = float(Ht + Wt) / float(Ho + Wo)
        return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
    else:
        return np.stack([smart_resize_k(x[:, :, i], fx, fy) for i in range(Co)], axis=2)


def padRightDownCorner(img, stride, padValue):
    h = img.shape[0]
    w = img.shape[1]

    pad = 4 * [None]
    pad[0] = 0 # up
    pad[1] = 0 # left
    pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
    pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right

    img_padded = img
    pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1))
    img_padded = np.concatenate((pad_up, img_padded), axis=0)
    pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1))
    img_padded = np.concatenate((pad_left, img_padded), axis=1)
    pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1))
    img_padded = np.concatenate((img_padded, pad_down), axis=0)
    pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1))
    img_padded = np.concatenate((img_padded, pad_right), axis=1)

    return img_padded, pad


def transfer(model, model_weights):
    transfered_model_weights = {}
    for weights_name in model.state_dict().keys():
        transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])]
    return transfered_model_weights


def draw_bodypose(canvas, candidate, subset):
    H, W, C = canvas.shape
    candidate = np.array(candidate)
    subset = np.array(subset)

    stickwidth = 4

    limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
               [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
               [1, 16], [16, 18], [3, 17], [6, 18]]

    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]]

    for i in range(17):
        for n in range(len(subset)):
            index = subset[n][np.array(limbSeq[i]) - 1]
            if -1 in index:
                continue
            Y = candidate[index.astype(int), 0] * float(W)
            X = candidate[index.astype(int), 1] * float(H)
            mX = np.mean(X)
            mY = np.mean(Y)
            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
            cv2.fillConvexPoly(canvas, polygon, colors[i])

    canvas = (canvas * 0.6).astype(np.uint8)

    for i in range(18):
        for n in range(len(subset)):
            index = int(subset[n][i])
            if index == -1:
                continue
            x, y = candidate[index][0:2]
            x = int(x * W)
            y = int(y * H)
            cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)

    return canvas

def alpha_blend_color(color, alpha):
    """blend color according to point conf
    """
    return [int(c * alpha) for c in color]

def draw_body_and_foot(canvas, candidate, subset, score, stick_width=4, draw_body=True, draw_feet=True, body_keypoint_size=4, draw_head=True):
    H, W, C = canvas.shape
    candidate = np.array(candidate)
    subset = np.array(subset)
    limbSeq_and_colors = []

    if draw_body:
        limbSeq_and_colors = [
            ([2, 3], [255, 0, 0]),   # Neck to Right Shoulder
            ([2, 6], [255, 85, 0]),  # Neck to Left Shoulder
            ([3, 4], [255, 170, 0]), # Right Shoulder to Right Elbow
            ([4, 5], [255, 255, 0]), # Right Elbow to Right Wrist
            ([6, 7], [170, 255, 0]), # Left Shoulder to Left Elbow
            ([7, 8], [85, 255, 0]),  # Left Elbow to Left Wrist
            ([2, 9], [0, 255, 0]),   # Neck to Right Hip
            ([9, 10], [0, 255, 85]), # Right Hip to Right Knee
            ([10, 11], [0, 255, 170]), # Right Knee to Right Ankle
            ([2, 12], [0, 255, 255]), # Neck to Left Hip
            ([12, 13], [0, 170, 255]), # Left Hip to Left Knee
            ([13, 14], [0, 85, 255]), # Left Knee to Left Ankle
        ]
    else:
        limbSeq_and_colors = [
            ([2, 3], [0, 0, 0]),   # Neck to Right Shoulder
            ([2, 6], [0, 0, 0]),  # Neck to Left Shoulder
            ([3, 4], [0, 0, 0]), # Right Shoulder to Right Elbow
            ([4, 5], [0, 0, 0]), # Right Elbow to Right Wrist
            ([6, 7], [0, 0, 0]), # Left Shoulder to Left Elbow
            ([7, 8], [0, 0, 0]),  # Left Elbow to Left Wrist
            ([2, 9], [0, 0, 0]),   # Neck to Right Hip
            ([9, 10], [0, 0, 0]), # Right Hip to Right Knee
            ([10, 11], [0, 0, 0]), # Right Knee to Right Ankle
            ([2, 12], [0, 0, 0]), # Neck to Left Hip
            ([12, 13], [0, 0, 0]), # Left Hip to Left Knee
            ([13, 14], [0, 0, 0]), # Left Knee to Left Ankle
        ]

    # Conditionally add head-related elements
    if draw_head:
        head_elements = [
            ([2, 1], [0, 0, 255]), # Neck to Nose
            ([1, 15], [0, 0, 255]), # Nose to Right Eye
            ([15, 17], [85, 0, 255]), # Right Eye to Right Ear
            ([1, 16], [170, 0, 255]), # Nose to Left Eye
            ([16, 18], [255, 0, 255]), # Left Eye to Left Ear
            ([3, 17], [255, 0, 170]), # Right Shoulder to Right Ear
            ([6, 18], [255, 0, 85])   # Left Shoulder to Left Ear
        ]
    else:
        head_elements = [
            ([2, 1], [0, 0, 0]),   # Neck to Nose
            ([1, 15], [0, 0, 0]), # Nose to Right Eye
            ([15, 17], [0, 0, 0]), # Right Eye to Right Ear
            ([1, 16], [0, 0, 0]), # Nose to Left Eye
            ([16, 18], [0, 0, 0]), # Left Eye to Left Ear
            ([3, 17], [0, 0, 0]), # Right Shoulder to Right Ear
            ([6, 18], [0, 0, 0])   # Left Shoulder to Left Ear
        ]
    if draw_feet:
        limbSeq_and_colors += [
            ([14, 19], [170, 255, 255]), # Left Ankle to Right Foot
            ([11, 20], [255, 255, 0]), # Right Ankle to Left Foot
        ]

    # Append head elements based on the condition
    limbSeq_and_colors += head_elements

    for limb_info in limbSeq_and_colors[:19]:
        limbSeq, color = limb_info
        for n in range(len(subset)):
            index = subset[n][np.array(limbSeq) - 1]
            if index[0] < 0.3 or index[1] < 0.3:
                continue
            Y = candidate[index.astype(int), 0] * float(W)
            X = candidate[index.astype(int), 1] * float(H)
            mX = np.mean(X)
            mY = np.mean(Y)
            length = np.sqrt((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2)
            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stick_width), int(angle), 0, 360, 1)
            cv2.fillConvexPoly(canvas, polygon, alpha_blend_color(color, index[0] * index[1]))

    canvas = (canvas * 0.6).astype(np.uint8)

    for limb_info in limbSeq_and_colors[:19]:
        limbSeq, color = limb_info
        for i in limbSeq:
            for n in range(len(subset)):
                index = int(subset[n][i - 1])
                if index == -1:
                    continue
                x, y = candidate[index][0:2]
                conf = score[n][i - 1]
                if not np.isfinite(x) or not np.isfinite(y):
                    continue
                x = int(np.clip(x * W, 0, W - 1))
                y = int(np.clip(y * H, 0, H - 1))
                cv2.circle(canvas, (x, y), body_keypoint_size, alpha_blend_color(color, conf), thickness=-1)

    return canvas


def draw_handpose(canvas, all_hand_peaks, draw_hands=True, hand_keypoint_size=4):
    H, W, C = canvas.shape

    edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
             [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]

    for peaks in all_hand_peaks:
        peaks = np.array(peaks)

        if draw_hands:
            for ie, e in enumerate(edges):
                x1, y1 = peaks[e[0]]
                x2, y2 = peaks[e[1]]
                x1 = int(x1 * W)
                y1 = int(y1 * H)
                x2 = int(x2 * W)
                y2 = int(y2 * H)
                if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
                    h = (ie / float(len(edges))) % 1.0
                    s, v = 1.0, 1.0
                    r, g, b = colorsys.hsv_to_rgb(h, s, v)
                    color = (int(255 * r), int(255 * g), int(255 * b))
                    cv2.line(canvas, (x1, y1), (x2, y2), color, thickness=2)
        if hand_keypoint_size > 0:
            for i, keypoint in enumerate(peaks):
                x, y = keypoint
                x = int(x * W)
                y = int(y * H)
                if x > eps and y > eps:
                    cv2.circle(canvas, (x, y), hand_keypoint_size, (0, 0, 255), thickness=-1)
    return canvas


def draw_facepose(canvas, all_lmks):
    H, W, C = canvas.shape
    for lmks in all_lmks:
        lmks = np.array(lmks)
        for lmk in lmks:
            x, y = lmk
            x = int(x * W)
            y = int(y * H)
            if x > eps and y > eps:
                cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
    return canvas


# detect hand according to body pose keypoints
# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
def handDetect(candidate, subset, oriImg):
    # right hand: wrist 4, elbow 3, shoulder 2
    # left hand: wrist 7, elbow 6, shoulder 5
    ratioWristElbow = 0.33
    detect_result = []
    image_height, image_width = oriImg.shape[0:2]
    for person in subset.astype(int):
        # if any of three not detected
        has_left = np.sum(person[[5, 6, 7]] == -1) == 0
        has_right = np.sum(person[[2, 3, 4]] == -1) == 0
        if not (has_left or has_right):
            continue
        hands = []
        #left hand
        if has_left:
            left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]]
            x1, y1 = candidate[left_shoulder_index][:2]
            x2, y2 = candidate[left_elbow_index][:2]
            x3, y3 = candidate[left_wrist_index][:2]
            hands.append([x1, y1, x2, y2, x3, y3, True])
        # right hand
        if has_right:
            right_shoulder_index, right_elbow_index, right_wrist_index = person[[2, 3, 4]]
            x1, y1 = candidate[right_shoulder_index][:2]
            x2, y2 = candidate[right_elbow_index][:2]
            x3, y3 = candidate[right_wrist_index][:2]
            hands.append([x1, y1, x2, y2, x3, y3, False])

        for x1, y1, x2, y2, x3, y3, is_left in hands:
            
            x = x3 + ratioWristElbow * (x3 - x2)
            y = y3 + ratioWristElbow * (y3 - y2)
            distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
            distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
            width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
            # x-y refers to the center --> offset to topLeft point
            # handRectangle.x -= handRectangle.width / 2.f;
            # handRectangle.y -= handRectangle.height / 2.f;
            x -= width / 2
            y -= width / 2  # width = height
            # overflow the image
            if x < 0: x = 0
            if y < 0: y = 0
            width1 = width
            width2 = width
            if x + width > image_width: width1 = image_width - x
            if y + width > image_height: width2 = image_height - y
            width = min(width1, width2)
            # the max hand box value is 20 pixels
            if width >= 20:
                detect_result.append([int(x), int(y), int(width), is_left])

    '''
    return value: [[x, y, w, True if left hand else False]].
    width=height since the network require squared input.
    x, y is the coordinate of top left 
    '''
    return detect_result


# Written by Lvmin
def faceDetect(candidate, subset, oriImg):
    # left right eye ear 14 15 16 17
    detect_result = []
    image_height, image_width = oriImg.shape[0:2]
    for person in subset.astype(int):
        has_head = person[0] > -1
        if not has_head:
            continue

        has_left_eye = person[14] > -1
        has_right_eye = person[15] > -1
        has_left_ear = person[16] > -1
        has_right_ear = person[17] > -1

        if not (has_left_eye or has_right_eye or has_left_ear or has_right_ear):
            continue

        head, left_eye, right_eye, left_ear, right_ear = person[[0, 14, 15, 16, 17]]

        width = 0.0
        x0, y0 = candidate[head][:2]

        if has_left_eye:
            x1, y1 = candidate[left_eye][:2]
            d = max(abs(x0 - x1), abs(y0 - y1))
            width = max(width, d * 3.0)

        if has_right_eye:
            x1, y1 = candidate[right_eye][:2]
            d = max(abs(x0 - x1), abs(y0 - y1))
            width = max(width, d * 3.0)

        if has_left_ear:
            x1, y1 = candidate[left_ear][:2]
            d = max(abs(x0 - x1), abs(y0 - y1))
            width = max(width, d * 1.5)

        if has_right_ear:
            x1, y1 = candidate[right_ear][:2]
            d = max(abs(x0 - x1), abs(y0 - y1))
            width = max(width, d * 1.5)

        x, y = x0, y0

        x -= width
        y -= width

        if x < 0:
            x = 0

        if y < 0:
            y = 0

        width1 = width * 2
        width2 = width * 2

        if x + width > image_width:
            width1 = image_width - x

        if y + width > image_height:
            width2 = image_height - y

        width = min(width1, width2)

        if width >= 20:
            detect_result.append([int(x), int(y), int(width)])

    return detect_result


# get max index of 2d array
def npmax(array):
    arrayindex = array.argmax(1)
    arrayvalue = array.max(1)
    i = arrayvalue.argmax()
    j = arrayindex[i]
    return i, j