processor.py

import os
import math
import base64
from typing import Callable

import cv2
import numpy as np
import matplotlib
from matplotlib import pyplot as plt
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
matplotlib.use('Agg')


class HandDetect:
    """
    Detect hand pose information from body pose estimation result.

    Args:
        ratioWristElbow(float): Ratio to adjust the wrist center, ,default is 0.33.
    """

    def __init__(self, ratioWristElbow: float = 0.33):
        self.ratioWristElbow = ratioWristElbow

    def __call__(self, candidate: np.ndarray, subset: np.ndarray, oriImg: np.ndarray):
        detect_result = []
        image_height, image_width = oriImg.shape[0:2]
        for person in subset.astype(int):
            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 + self.ratioWristElbow * (x3 - x2)
                y = y3 + self.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 -= width / 2
                y -= width / 2

                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)

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

        return detect_result


class PadDownRight:
    """
    Get padding images.

    Args:
        stride(int): Stride for calculate pad value for edges.
        padValue(int): Initialization for new area.
    """

    def __init__(self, stride: int = 8, padValue: int = 128):
        self.stride = stride
        self.padValue = padValue

    def __call__(self, img: np.ndarray):
        h, w = img.shape[0:2]
        pad = 4 * [0]
        pad[2] = 0 if (h % self.stride == 0) else self.stride - (h % self.stride)  # down
        pad[3] = 0 if (w % self.stride == 0) else self.stride - (w % self.stride)  # right

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

        return img_padded, pad


class RemovePadding:
    """
    Remove the padding values.

    Args:
        stride(int): Scales for resizing the images.

    """

    def __init__(self, stride: int = 8):
        self.stride = stride

    def __call__(self, data: np.ndarray, imageToTest_padded: np.ndarray, oriImg: np.ndarray, pad: list) -> np.ndarray:
        heatmap = np.transpose(np.squeeze(data), (1, 2, 0))
        heatmap = cv2.resize(heatmap, (0, 0), fx=self.stride, fy=self.stride, interpolation=cv2.INTER_CUBIC)
        heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
        heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)

        return heatmap


class DrawPose:
    """
    Draw Pose estimation results on canvas.

    Args:
        stickwidth(int): Angle value to draw approximate ellipse curve, default is 4.

    """

    def __init__(self, stickwidth: int = 4):
        self.stickwidth = stickwidth

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

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

    def __call__(self, canvas: np.ndarray, candidate: np.ndarray, subset: np.ndarray):
        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]
                cv2.circle(canvas, (int(x), int(y)), 4, self.colors[i], thickness=-1)
        for i in range(17):
            for n in range(len(subset)):
                index = subset[n][np.array(self.limbSeq[i]) - 1]
                if -1 in index:
                    continue
                cur_canvas = canvas.copy()
                Y = candidate[index.astype(int), 0]
                X = candidate[index.astype(int), 1]
                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), self.stickwidth), int(angle), 0, 360,
                                           1)
                cv2.fillConvexPoly(cur_canvas, polygon, self.colors[i])
                canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
        return canvas


class DrawHandPose:
    """
        Draw hand pose estimation results on canvas.
        Args:
            show_number(bool): Whether to show estimation ids in canvas, default is False.

        """

    def __init__(self, show_number: bool = False):
        self.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]]
        self.show_number = show_number

    def __call__(self, canvas: np.ndarray, all_hand_peaks: list):
        fig = Figure(figsize=plt.figaspect(canvas))

        fig.subplots_adjust(0, 0, 1, 1)
        fig.subplots_adjust(bottom=0, top=1, left=0, right=1)
        bg = FigureCanvas(fig)
        ax = fig.subplots()
        ax.axis('off')
        ax.imshow(canvas)

        width, height = ax.figure.get_size_inches() * ax.figure.get_dpi()

        for peaks in all_hand_peaks:
            for ie, e in enumerate(self.edges):
                if np.sum(np.all(peaks[e], axis=1) == 0) == 0:
                    x1, y1 = peaks[e[0]]
                    x2, y2 = peaks[e[1]]
                    ax.plot([x1, x2], [y1, y2],
                            color=matplotlib.colors.hsv_to_rgb([ie / float(len(self.edges)), 1.0, 1.0]))

            for i, keyponit in enumerate(peaks):
                x, y = keyponit
                ax.plot(x, y, 'r.')
                if self.show_number:
                    ax.text(x, y, str(i))
        bg.draw()
        canvas = np.frombuffer(bg.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3)
        return canvas


class ResizeScaling:
    """Resize images by scaling method.

    Args:
        target(int): Target image size.
        interpolation(Callable): Interpolation method.
    """

    def __init__(self, target: int = 368, interpolation: Callable = cv2.INTER_CUBIC):
        self.target = target
        self.interpolation = interpolation

    def __call__(self, img, scale_search):
        scale = scale_search * self.target / img.shape[0]
        resize_img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=self.interpolation)
        return resize_img


def npmax(array: np.ndarray):
    """Get max value and index."""
    arrayindex = array.argmax(1)
    arrayvalue = array.max(1)
    i = arrayvalue.argmax()
    j = arrayindex[i]
    return i, j


def cv2_to_base64(image: np.ndarray):
    data = cv2.imencode('.jpg', image)[1]
    return base64.b64encode(data.tostring()).decode('utf8')


def base64_to_cv2(b64str: str):
    data = base64.b64decode(b64str.encode('utf8'))
    data = np.fromstring(data, np.uint8)
    data = cv2.imdecode(data, cv2.IMREAD_COLOR)
    return data