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