Spaces:
Configuration error
Configuration error
| """ | |
| Copyright (c) Microsoft Corporation. | |
| Licensed under the MIT license. | |
| """ | |
| import cv2 | |
| import math | |
| import json | |
| from PIL import Image | |
| import os.path as op | |
| import numpy as np | |
| import code | |
| from custom_mesh_graphormer.utils.tsv_file import TSVFile, CompositeTSVFile | |
| from custom_mesh_graphormer.utils.tsv_file_ops import load_linelist_file, load_from_yaml_file, find_file_path_in_yaml | |
| from custom_mesh_graphormer.utils.image_ops import img_from_base64, crop, flip_img, flip_pose, flip_kp, transform, rot_aa | |
| import torch | |
| import torchvision.transforms as transforms | |
| class MeshTSVDataset(object): | |
| def __init__(self, img_file, label_file=None, hw_file=None, | |
| linelist_file=None, is_train=True, cv2_output=False, scale_factor=1): | |
| self.img_file = img_file | |
| self.label_file = label_file | |
| self.hw_file = hw_file | |
| self.linelist_file = linelist_file | |
| self.img_tsv = self.get_tsv_file(img_file) | |
| self.label_tsv = None if label_file is None else self.get_tsv_file(label_file) | |
| self.hw_tsv = None if hw_file is None else self.get_tsv_file(hw_file) | |
| if self.is_composite: | |
| assert op.isfile(self.linelist_file) | |
| self.line_list = [i for i in range(self.hw_tsv.num_rows())] | |
| else: | |
| self.line_list = load_linelist_file(linelist_file) | |
| self.cv2_output = cv2_output | |
| self.normalize_img = transforms.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]) | |
| self.is_train = is_train | |
| self.scale_factor = 0.25 # rescale bounding boxes by a factor of [1-options.scale_factor,1+options.scale_factor] | |
| self.noise_factor = 0.4 | |
| self.rot_factor = 30 # Random rotation in the range [-rot_factor, rot_factor] | |
| self.img_res = 224 | |
| self.image_keys = self.prepare_image_keys() | |
| self.joints_definition = ('R_Ankle', 'R_Knee', 'R_Hip', 'L_Hip', 'L_Knee', 'L_Ankle', 'R_Wrist', 'R_Elbow', 'R_Shoulder', 'L_Shoulder', | |
| 'L_Elbow','L_Wrist','Neck','Top_of_Head','Pelvis','Thorax','Spine','Jaw','Head','Nose','L_Eye','R_Eye','L_Ear','R_Ear') | |
| self.pelvis_index = self.joints_definition.index('Pelvis') | |
| def get_tsv_file(self, tsv_file): | |
| if tsv_file: | |
| if self.is_composite: | |
| return CompositeTSVFile(tsv_file, self.linelist_file, | |
| root=self.root) | |
| tsv_path = find_file_path_in_yaml(tsv_file, self.root) | |
| return TSVFile(tsv_path) | |
| def get_valid_tsv(self): | |
| # sorted by file size | |
| if self.hw_tsv: | |
| return self.hw_tsv | |
| if self.label_tsv: | |
| return self.label_tsv | |
| def prepare_image_keys(self): | |
| tsv = self.get_valid_tsv() | |
| return [tsv.get_key(i) for i in range(tsv.num_rows())] | |
| def prepare_image_key_to_index(self): | |
| tsv = self.get_valid_tsv() | |
| return {tsv.get_key(i) : i for i in range(tsv.num_rows())} | |
| def augm_params(self): | |
| """Get augmentation parameters.""" | |
| flip = 0 # flipping | |
| pn = np.ones(3) # per channel pixel-noise | |
| rot = 0 # rotation | |
| sc = 1 # scaling | |
| if self.is_train: | |
| # We flip with probability 1/2 | |
| if np.random.uniform() <= 0.5: | |
| flip = 1 | |
| # Each channel is multiplied with a number | |
| # in the area [1-opt.noiseFactor,1+opt.noiseFactor] | |
| pn = np.random.uniform(1-self.noise_factor, 1+self.noise_factor, 3) | |
| # The rotation is a number in the area [-2*rotFactor, 2*rotFactor] | |
| rot = min(2*self.rot_factor, | |
| max(-2*self.rot_factor, np.random.randn()*self.rot_factor)) | |
| # The scale is multiplied with a number | |
| # in the area [1-scaleFactor,1+scaleFactor] | |
| sc = min(1+self.scale_factor, | |
| max(1-self.scale_factor, np.random.randn()*self.scale_factor+1)) | |
| # but it is zero with probability 3/5 | |
| if np.random.uniform() <= 0.6: | |
| rot = 0 | |
| return flip, pn, rot, sc | |
| def rgb_processing(self, rgb_img, center, scale, rot, flip, pn): | |
| """Process rgb image and do augmentation.""" | |
| rgb_img = crop(rgb_img, center, scale, | |
| [self.img_res, self.img_res], rot=rot) | |
| # flip the image | |
| if flip: | |
| rgb_img = flip_img(rgb_img) | |
| # in the rgb image we add pixel noise in a channel-wise manner | |
| rgb_img[:,:,0] = np.minimum(255.0, np.maximum(0.0, rgb_img[:,:,0]*pn[0])) | |
| rgb_img[:,:,1] = np.minimum(255.0, np.maximum(0.0, rgb_img[:,:,1]*pn[1])) | |
| rgb_img[:,:,2] = np.minimum(255.0, np.maximum(0.0, rgb_img[:,:,2]*pn[2])) | |
| # (3,224,224),float,[0,1] | |
| rgb_img = np.transpose(rgb_img.astype('float32'),(2,0,1))/255.0 | |
| return rgb_img | |
| def j2d_processing(self, kp, center, scale, r, f): | |
| """Process gt 2D keypoints and apply all augmentation transforms.""" | |
| nparts = kp.shape[0] | |
| for i in range(nparts): | |
| kp[i,0:2] = transform(kp[i,0:2]+1, center, scale, | |
| [self.img_res, self.img_res], rot=r) | |
| # convert to normalized coordinates | |
| kp[:,:-1] = 2.*kp[:,:-1]/self.img_res - 1. | |
| # flip the x coordinates | |
| if f: | |
| kp = flip_kp(kp) | |
| kp = kp.astype('float32') | |
| return kp | |
| def j3d_processing(self, S, r, f): | |
| """Process gt 3D keypoints and apply all augmentation transforms.""" | |
| # in-plane rotation | |
| rot_mat = np.eye(3) | |
| if not r == 0: | |
| rot_rad = -r * np.pi / 180 | |
| sn,cs = np.sin(rot_rad), np.cos(rot_rad) | |
| rot_mat[0,:2] = [cs, -sn] | |
| rot_mat[1,:2] = [sn, cs] | |
| S[:, :-1] = np.einsum('ij,kj->ki', rot_mat, S[:, :-1]) | |
| # flip the x coordinates | |
| if f: | |
| S = flip_kp(S) | |
| S = S.astype('float32') | |
| return S | |
| def pose_processing(self, pose, r, f): | |
| """Process SMPL theta parameters and apply all augmentation transforms.""" | |
| # rotation or the pose parameters | |
| pose = pose.astype('float32') | |
| pose[:3] = rot_aa(pose[:3], r) | |
| # flip the pose parameters | |
| if f: | |
| pose = flip_pose(pose) | |
| # (72),float | |
| pose = pose.astype('float32') | |
| return pose | |
| def get_line_no(self, idx): | |
| return idx if self.line_list is None else self.line_list[idx] | |
| def get_image(self, idx): | |
| line_no = self.get_line_no(idx) | |
| row = self.img_tsv[line_no] | |
| # use -1 to support old format with multiple columns. | |
| cv2_im = img_from_base64(row[-1]) | |
| if self.cv2_output: | |
| return cv2_im.astype(np.float32, copy=True) | |
| cv2_im = cv2.cvtColor(cv2_im, cv2.COLOR_BGR2RGB) | |
| return cv2_im | |
| def get_annotations(self, idx): | |
| line_no = self.get_line_no(idx) | |
| if self.label_tsv is not None: | |
| row = self.label_tsv[line_no] | |
| annotations = json.loads(row[1]) | |
| return annotations | |
| else: | |
| return [] | |
| def get_target_from_annotations(self, annotations, img_size, idx): | |
| # This function will be overwritten by each dataset to | |
| # decode the labels to specific formats for each task. | |
| return annotations | |
| def get_img_info(self, idx): | |
| if self.hw_tsv is not None: | |
| line_no = self.get_line_no(idx) | |
| row = self.hw_tsv[line_no] | |
| try: | |
| # json string format with "height" and "width" being the keys | |
| return json.loads(row[1])[0] | |
| except ValueError: | |
| # list of strings representing height and width in order | |
| hw_str = row[1].split(' ') | |
| hw_dict = {"height": int(hw_str[0]), "width": int(hw_str[1])} | |
| return hw_dict | |
| def get_img_key(self, idx): | |
| line_no = self.get_line_no(idx) | |
| # based on the overhead of reading each row. | |
| if self.hw_tsv: | |
| return self.hw_tsv[line_no][0] | |
| elif self.label_tsv: | |
| return self.label_tsv[line_no][0] | |
| else: | |
| return self.img_tsv[line_no][0] | |
| def __len__(self): | |
| if self.line_list is None: | |
| return self.img_tsv.num_rows() | |
| else: | |
| return len(self.line_list) | |
| def __getitem__(self, idx): | |
| img = self.get_image(idx) | |
| img_key = self.get_img_key(idx) | |
| annotations = self.get_annotations(idx) | |
| annotations = annotations[0] | |
| center = annotations['center'] | |
| scale = annotations['scale'] | |
| has_2d_joints = annotations['has_2d_joints'] | |
| has_3d_joints = annotations['has_3d_joints'] | |
| joints_2d = np.asarray(annotations['2d_joints']) | |
| joints_3d = np.asarray(annotations['3d_joints']) | |
| if joints_2d.ndim==3: | |
| joints_2d = joints_2d[0] | |
| if joints_3d.ndim==3: | |
| joints_3d = joints_3d[0] | |
| # Get SMPL parameters, if available | |
| has_smpl = np.asarray(annotations['has_smpl']) | |
| pose = np.asarray(annotations['pose']) | |
| betas = np.asarray(annotations['betas']) | |
| try: | |
| gender = annotations['gender'] | |
| except KeyError: | |
| gender = 'none' | |
| # Get augmentation parameters | |
| flip,pn,rot,sc = self.augm_params() | |
| # Process image | |
| img = self.rgb_processing(img, center, sc*scale, rot, flip, pn) | |
| img = torch.from_numpy(img).float() | |
| # Store image before normalization to use it in visualization | |
| transfromed_img = self.normalize_img(img) | |
| # normalize 3d pose by aligning the pelvis as the root (at origin) | |
| root_pelvis = joints_3d[self.pelvis_index,:-1] | |
| joints_3d[:,:-1] = joints_3d[:,:-1] - root_pelvis[None,:] | |
| # 3d pose augmentation (random flip + rotation, consistent to image and SMPL) | |
| joints_3d_transformed = self.j3d_processing(joints_3d.copy(), rot, flip) | |
| # 2d pose augmentation | |
| joints_2d_transformed = self.j2d_processing(joints_2d.copy(), center, sc*scale, rot, flip) | |
| ################################### | |
| # Masking percantage | |
| # We observe that 30% works better for human body mesh. Further details are reported in the paper. | |
| mvm_percent = 0.3 | |
| ################################### | |
| mjm_mask = np.ones((14,1)) | |
| if self.is_train: | |
| num_joints = 14 | |
| pb = np.random.random_sample() | |
| masked_num = int(pb * mvm_percent * num_joints) # at most x% of the joints could be masked | |
| indices = np.random.choice(np.arange(num_joints),replace=False,size=masked_num) | |
| mjm_mask[indices,:] = 0.0 | |
| mjm_mask = torch.from_numpy(mjm_mask).float() | |
| mvm_mask = np.ones((431,1)) | |
| if self.is_train: | |
| num_vertices = 431 | |
| pb = np.random.random_sample() | |
| masked_num = int(pb * mvm_percent * num_vertices) # at most x% of the vertices could be masked | |
| indices = np.random.choice(np.arange(num_vertices),replace=False,size=masked_num) | |
| mvm_mask[indices,:] = 0.0 | |
| mvm_mask = torch.from_numpy(mvm_mask).float() | |
| meta_data = {} | |
| meta_data['ori_img'] = img | |
| meta_data['pose'] = torch.from_numpy(self.pose_processing(pose, rot, flip)).float() | |
| meta_data['betas'] = torch.from_numpy(betas).float() | |
| meta_data['joints_3d'] = torch.from_numpy(joints_3d_transformed).float() | |
| meta_data['has_3d_joints'] = has_3d_joints | |
| meta_data['has_smpl'] = has_smpl | |
| meta_data['mjm_mask'] = mjm_mask | |
| meta_data['mvm_mask'] = mvm_mask | |
| # Get 2D keypoints and apply augmentation transforms | |
| meta_data['has_2d_joints'] = has_2d_joints | |
| meta_data['joints_2d'] = torch.from_numpy(joints_2d_transformed).float() | |
| meta_data['scale'] = float(sc * scale) | |
| meta_data['center'] = np.asarray(center).astype(np.float32) | |
| meta_data['gender'] = gender | |
| return img_key, transfromed_img, meta_data | |
| class MeshTSVYamlDataset(MeshTSVDataset): | |
| """ TSVDataset taking a Yaml file for easy function call | |
| """ | |
| def __init__(self, yaml_file, is_train=True, cv2_output=False, scale_factor=1): | |
| self.cfg = load_from_yaml_file(yaml_file) | |
| self.is_composite = self.cfg.get('composite', False) | |
| self.root = op.dirname(yaml_file) | |
| if self.is_composite==False: | |
| img_file = find_file_path_in_yaml(self.cfg['img'], self.root) | |
| label_file = find_file_path_in_yaml(self.cfg.get('label', None), | |
| self.root) | |
| hw_file = find_file_path_in_yaml(self.cfg.get('hw', None), self.root) | |
| linelist_file = find_file_path_in_yaml(self.cfg.get('linelist', None), | |
| self.root) | |
| else: | |
| img_file = self.cfg['img'] | |
| hw_file = self.cfg['hw'] | |
| label_file = self.cfg.get('label', None) | |
| linelist_file = find_file_path_in_yaml(self.cfg.get('linelist', None), | |
| self.root) | |
| super(MeshTSVYamlDataset, self).__init__( | |
| img_file, label_file, hw_file, linelist_file, is_train, cv2_output=cv2_output, scale_factor=scale_factor) | |