NTED/demo_dataset.py

import os
import cv2
import math
import numpy as np
from PIL import Image

import torch
import torchvision.transforms.functional as F

class DemoDataset(object):
    def __init__(self):
        super().__init__()    
        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]]
        
        self.LIMBSEQ_hands = [[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], \
            [21, 22], [22, 23], [23, 24], [24, 25], \
            [21, 26], [26, 27], [27, 28], [28, 29], \
            [21, 30], [30, 31], [31, 32], [32, 33], \
            [21, 34], [34, 35], [35, 36], [36, 37], \
            [21, 38], [38, 39], [39, 40], [40, 41]]
        
        self.COLORS_hands = [[85, 0, 0], [170, 0, 0], [85, 85, 0], [85, 170, 0], [170, 85, 0], [170, 170, 0], [85, 85, 85], \
            [85, 85, 170], [85, 170, 85], [85, 170, 170], [0, 85, 0], [0, 170, 0], [0, 85, 85], [0, 85, 170], \
            [0, 170, 85], [0, 170, 170], [50, 0, 0], [135, 0, 0], [50, 50, 0], [50, 135, 0], [135, 50, 0], \
            [135, 135, 0], [50, 50, 50], [50, 50, 135], [50, 135, 50], [50, 135, 135], [0, 50, 0], [0, 135, 0], \
            [0, 50, 50], [0, 50, 135], [0, 135, 50], [0, 135, 135], [100, 0, 0], [200, 0, 0], [100, 100, 0], \
            [100, 200, 0], [200, 100, 0], [200, 200, 0], [100, 100, 100], [100, 100, 200], [100, 200, 100], [100, 200, 200]
            ]
        
        self.img_size = tuple([512, 352])
    
    def load_item(self, img, pose, handpose=None):

        reference_img = self.get_image_tensor(img)[None,:]
        label, ske = self.get_label_tensor(pose, handpose)
        label = label[None,:]

        return {'reference_image':reference_img, 'target_skeleton':label, 'skeleton_img': ske}
    
    def get_image_tensor(self, bgr_img):
        img = Image.fromarray(cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB))
        img = F.resize(img, self.img_size)
        img = F.to_tensor(img)
        img = F.normalize(img, (0.5, 0.5, 0.5),(0.5, 0.5, 0.5))
        return img    

    def get_label_tensor(self, pose, hand_pose=None):
        canvas = np.zeros((self.img_size[0], self.img_size[1], 3)).astype(np.uint8)
        keypoint = np.array(pose)
        if hand_pose is not None:
            keypoint_hands = np.array(hand_pose)
        else:
            keypoint_hands = None
        
        # keypoint = self.trans_keypoins(keypoint)
        
        stickwidth = 4
        for i in range(18):
            x, y = keypoint[i, 0:2]
            if x == -1 or y == -1:
                continue
            cv2.circle(canvas, (int(x), int(y)), 4, self.COLORS[i], thickness=-1)
        if keypoint_hands is not None:
            for i in range(42):
                    x, y = keypoint_hands[i, 0:2]
                    if x == -1 or y == -1:
                        continue
                    cv2.circle(canvas, (int(x), int(y)), 4, self.COLORS_hands[i], thickness=-1)
        
        joints = []
        for i in range(17):
            Y = keypoint[np.array(self.LIMBSEQ[i])-1, 0]
            X = keypoint[np.array(self.LIMBSEQ[i])-1, 1]            
            cur_canvas = canvas.copy()
            if -1 in Y or -1 in X:
                joints.append(np.zeros_like(cur_canvas[:, :, 0]))
                continue
            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(cur_canvas, polygon, self.COLORS[i])
            canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)

            joint = np.zeros_like(cur_canvas[:, :, 0])
            cv2.fillConvexPoly(joint, polygon, 255)
            joint = cv2.addWeighted(joint, 0.4, joint, 0.6, 0)
            joints.append(joint)
        if keypoint_hands is not None:
            for i in range(40):
                    Y = keypoint_hands[np.array(self.LIMBSEQ_hands[i]), 0]
                    X = keypoint_hands[np.array(self.LIMBSEQ_hands[i]), 1]            
                    cur_canvas = canvas.copy()
                    if -1 in Y or -1 in X:
                        if (i+1) % 4 == 0:
                            joints.append(np.zeros_like(cur_canvas[:, :, 0]))
                        continue
                    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), int(stickwidth/2)), int(angle), 0, 360, 1)
                    cv2.fillConvexPoly(cur_canvas, polygon, self.COLORS_hands[i])
                    canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
                    
                    # 一根手指一个通道
                    if i % 4 == 0:
                        joint = np.zeros_like(cur_canvas[:, :, 0])
                    cv2.fillConvexPoly(joint, polygon, 255)
                    joint = cv2.addWeighted(joint, 0.4, joint, 0.6, 0)
                    if (i+1) % 4 == 0:
                        joints.append(joint)
        
        pose = F.to_tensor(Image.fromarray(cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)))
        
        tensors_dist = 0
        e = 1
        for i in range(len(joints)):
            im_dist = cv2.distanceTransform(255-joints[i], cv2.DIST_L1, 3)
            im_dist = np.clip((im_dist / 3), 0, 255).astype(np.uint8)
            tensor_dist = F.to_tensor(Image.fromarray(im_dist))
            tensors_dist = tensor_dist if e == 1 else torch.cat([tensors_dist, tensor_dist])
            e += 1

        label_tensor = torch.cat((pose, tensors_dist), dim=0)
            
        return label_tensor, canvas     
    
    def tensor2im(self, image_tensor, imtype=np.uint8, normalize=True,
              three_channel_output=True):
        r"""Convert tensor to image.

        Args:
            image_tensor (torch.tensor or list of torch.tensor): If tensor then
                (NxCxHxW) or (NxTxCxHxW) or (CxHxW).
            imtype (np.dtype): Type of output image.
            normalize (bool): Is the input image normalized or not?
                three_channel_output (bool): Should single channel images be made 3
                channel in output?

        Returns:
            (numpy.ndarray, list if case 1, 2 above).
        """
        if image_tensor is None:
            return None
        if isinstance(image_tensor, list):
            return [self.tensor2im(x, imtype, normalize) for x in image_tensor]
        if image_tensor.dim() == 5 or image_tensor.dim() == 4:
            return [self.tensor2im(image_tensor[idx], imtype, normalize)
                    for idx in range(image_tensor.size(0))]

        if image_tensor.dim() == 3:
            image_numpy = image_tensor.detach().float().numpy()
            if normalize:
                image_numpy = (np.transpose(
                    image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0
            else:
                image_numpy = np.transpose(image_numpy, (1, 2, 0)) * 255.0
            image_numpy = np.clip(image_numpy, 0, 255)
            if image_numpy.shape[2] == 1 and three_channel_output:
                image_numpy = np.repeat(image_numpy, 3, axis=2)
            elif image_numpy.shape[2] > 3:
                image_numpy = image_numpy[:, :, :3]
            return image_numpy.astype(imtype)

    def trans_keypoins(self, keypoints):
        missing_keypoint_index = keypoints == -1

        keypoints[missing_keypoint_index] = -1
        return keypoints