Spaces:
Configuration error
Configuration error
| import sys | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.utils.data import Dataset, DataLoader | |
| import numpy as np | |
| import os | |
| import random | |
| import copy | |
| import json | |
| from collections import defaultdict | |
| from lib.utils.utils_data import crop_scale, flip_data, resample, split_clips | |
| def posetrack2h36m(x): | |
| ''' | |
| Input: x (T x V x C) | |
| PoseTrack keypoints = [ 'nose', | |
| 'head_bottom', | |
| 'head_top', | |
| 'left_ear', | |
| 'right_ear', | |
| 'left_shoulder', | |
| 'right_shoulder', | |
| 'left_elbow', | |
| 'right_elbow', | |
| 'left_wrist', | |
| 'right_wrist', | |
| 'left_hip', | |
| 'right_hip', | |
| 'left_knee', | |
| 'right_knee', | |
| 'left_ankle', | |
| 'right_ankle'] | |
| H36M: | |
| 0: 'root', | |
| 1: 'rhip', | |
| 2: 'rkne', | |
| 3: 'rank', | |
| 4: 'lhip', | |
| 5: 'lkne', | |
| 6: 'lank', | |
| 7: 'belly', | |
| 8: 'neck', | |
| 9: 'nose', | |
| 10: 'head', | |
| 11: 'lsho', | |
| 12: 'lelb', | |
| 13: 'lwri', | |
| 14: 'rsho', | |
| 15: 'relb', | |
| 16: 'rwri' | |
| ''' | |
| y = np.zeros(x.shape) | |
| y[:,0,:] = (x[:,11,:] + x[:,12,:]) * 0.5 | |
| y[:,1,:] = x[:,12,:] | |
| y[:,2,:] = x[:,14,:] | |
| y[:,3,:] = x[:,16,:] | |
| y[:,4,:] = x[:,11,:] | |
| y[:,5,:] = x[:,13,:] | |
| y[:,6,:] = x[:,15,:] | |
| y[:,8,:] = x[:,1,:] | |
| y[:,7,:] = (y[:,0,:] + y[:,8,:]) * 0.5 | |
| y[:,9,:] = x[:,0,:] | |
| y[:,10,:] = x[:,2,:] | |
| y[:,11,:] = x[:,5,:] | |
| y[:,12,:] = x[:,7,:] | |
| y[:,13,:] = x[:,9,:] | |
| y[:,14,:] = x[:,6,:] | |
| y[:,15,:] = x[:,8,:] | |
| y[:,16,:] = x[:,10,:] | |
| y[:,0,2] = np.minimum(x[:,11,2], x[:,12,2]) | |
| y[:,7,2] = np.minimum(y[:,0,2], y[:,8,2]) | |
| return y | |
| class PoseTrackDataset2D(Dataset): | |
| def __init__(self, flip=True, scale_range=[0.25, 1]): | |
| super(PoseTrackDataset2D, self).__init__() | |
| self.flip = flip | |
| data_root = "data/motion2d/posetrack18_annotations/train/" | |
| file_list = sorted(os.listdir(data_root)) | |
| all_motions = [] | |
| all_motions_filtered = [] | |
| self.scale_range = scale_range | |
| for filename in file_list: | |
| with open(os.path.join(data_root, filename), 'r') as file: | |
| json_dict = json.load(file) | |
| annots = json_dict['annotations'] | |
| imgs = json_dict['images'] | |
| motions = defaultdict(list) | |
| for annot in annots: | |
| tid = annot['track_id'] | |
| pose2d = np.array(annot['keypoints']).reshape(-1,3) | |
| motions[tid].append(pose2d) | |
| all_motions += list(motions.values()) | |
| for motion in all_motions: | |
| if len(motion)<30: | |
| continue | |
| motion = np.array(motion[:30]) | |
| if np.sum(motion[:,:,2]) <= 306: # Valid joint num threshold | |
| continue | |
| motion = crop_scale(motion, self.scale_range) | |
| motion = posetrack2h36m(motion) | |
| motion[motion[:,:,2]==0] = 0 | |
| if np.sum(motion[:,0,2]) < 30: | |
| continue # Root all visible (needed for framewise rootrel) | |
| all_motions_filtered.append(motion) | |
| all_motions_filtered = np.array(all_motions_filtered) | |
| self.motions_2d = all_motions_filtered | |
| def __len__(self): | |
| 'Denotes the total number of samples' | |
| return len(self.motions_2d) | |
| def __getitem__(self, index): | |
| 'Generates one sample of data' | |
| motion_2d = torch.FloatTensor(self.motions_2d[index]) | |
| if self.flip and random.random()>0.5: | |
| motion_2d = flip_data(motion_2d) | |
| return motion_2d, motion_2d | |
| class InstaVDataset2D(Dataset): | |
| def __init__(self, n_frames=81, data_stride=27, flip=True, valid_threshold=0.0, scale_range=[0.25, 1]): | |
| super(InstaVDataset2D, self).__init__() | |
| self.flip = flip | |
| self.scale_range = scale_range | |
| motion_all = np.load('data/motion2d/InstaVariety/motion_all.npy') | |
| id_all = np.load('data/motion2d/InstaVariety/id_all.npy') | |
| split_id = split_clips(id_all, n_frames, data_stride) | |
| motions_2d = motion_all[split_id] # [N, T, 17, 3] | |
| valid_idx = (motions_2d[:,0,0,2] > valid_threshold) | |
| self.motions_2d = motions_2d[valid_idx] | |
| def __len__(self): | |
| 'Denotes the total number of samples' | |
| return len(self.motions_2d) | |
| def __getitem__(self, index): | |
| 'Generates one sample of data' | |
| motion_2d = self.motions_2d[index] | |
| motion_2d = crop_scale(motion_2d, self.scale_range) | |
| motion_2d[motion_2d[:,:,2]==0] = 0 | |
| if self.flip and random.random()>0.5: | |
| motion_2d = flip_data(motion_2d) | |
| motion_2d = torch.FloatTensor(motion_2d) | |
| return motion_2d, motion_2d | |