os.makedirs(to_file)
files = os.listdir(from_file)
for f in files:
if os.path.isdir(from_file + '\\' + f):
copydirs(from_file + '\\' + f, to_file + '\\' + f)
else:
if '.git' not in from_file:
shutil.copy(from_file + '\\' + f, to_file + '\\' + f)
copydirs('.\\', log_dir + '\\src')
## initilize the skeleton ##
skeleton_mocap = Skeleton(offsets=opt['data']['offsets'],
parents=opt['data']['parents'])
skeleton_mocap.cuda()
skeleton_mocap.remove_joints(opt['data']['joints_to_remove'])
## load train data ##
lafan_data_train = LaFan1(opt['data']['data_dir'], \
seq_len = opt['model']['seq_length'], \
offset = opt['data']['offset'],\
train = True, debug=opt['train']['debug'])
x_mean = lafan_data_train.x_mean.cuda()
x_std = lafan_data_train.x_std.cuda().view(1, 1,
opt['model']['num_joints'], 3)
if opt['train']['debug']:
opt['data']['num_workers'] = 1
lafan_loader_train = DataLoader(lafan_data_train, \
batch_size=opt['train']['batch_size'], \
shuffle=True, num_workers=opt['data']['num_workers'])
class Human36mDataset(Dataset):
def __init__(self,
data_root,
max_seq_len,
delta_len,
n_breakpoints=0,
speed_range=[1, 1],
acc_range=[-1, 1],
train=True,
remove_static_joints=True,
mode='train'):
# Set arguments
self.data_root = osp.abspath(osp.expanduser(data_root))
self.max_seq_len = max_seq_len
self.delta_len = delta_len
self.speed_range = speed_range
self.n_breakpoints = n_breakpoints
self.acc_range = acc_range
self.remove_static_joints = remove_static_joints
self.train = train
self.subseq_len = self.delta_len * 2
self.skeleton = Skeleton(
parents=[
-1, 0, 1, 2, 3, 4, 0, 6, 7, 8, 9, 0, 11, 12, 13, 14, 12, 16,
17, 18, 19, 20, 19, 22, 12, 24, 25, 26, 27, 28, 27, 30
],
joints_left=[6, 7, 8, 9, 10, 16, 17, 18, 19, 20, 21, 22, 23],
joints_right=[1, 2, 3, 4, 5, 24, 25, 26, 27, 28, 29, 30, 31])
assert mode in [
'train', 'test'
], 'Invalid mode for Human36mDataset. Must be \'train\' or \'test\'.'
self.mode = mode
# Read dataset
self.raw_data = read_human36m(self.data_root, self.mode)
# Reformat data
self.data = reformat_data(self.raw_data)
# Filter out sequence that is shorter than `max_seq_len`
self.data['pose']['3d'] = list(
filter(lambda x: x.shape[0] >= self.max_seq_len,
self.data['pose']['3d']))
self.data['pose']['2d'] = list(
filter(lambda x: x.shape[0] >= self.max_seq_len,
self.data['pose']['2d']))
# Remove static joints
if self.remove_static_joints:
# Bring the skeleton to 17 joints instead of the original 32
self.remove_joints(
[4, 5, 9, 10, 11, 16, 20, 21, 22, 23, 24, 28, 29, 30, 31])
# Rewire shoulders to the correct parents
self.skeleton._parents[11] = 8
self.skeleton._parents[14] = 8
# Preprocess data; NOTE: make sure index of pivot joint is correct, if used.
align_and_normalize_dataset_v2(self.data)
def __getitem__(self, idx):
# Fetch data
pose_2d = self.data['pose']['2d'][idx]
pose_3d = self.data['pose']['3d'][idx]
camera_view = self.data['camera_view'][idx]
# Crop a sub-sequence
total_len = pose_3d.shape[0]
start = np.random.randint(low=0,
high=total_len -
self.speed_range[1] * self.max_seq_len + 1)
pose_2d_cropped, pose_3d_cropped, speed = [], [], []
if self.n_breakpoints > 0: # Varying video speed
start_offset = 5
bps = [0, self.max_seq_len] + list(
np.random.randint(low=1 + start_offset,
high=self.max_seq_len - start_offset,
size=self.n_breakpoints))
bps = sorted(bps)
speed.append(
np.random.randint(low=self.speed_range[0],
high=self.speed_range[1] + 1))
for bp_i, bp in enumerate(bps[1:]):
end = start + (bp - bps[bp_i]) * speed[-1]
pose_2d_cropped.append(pose_2d[start:end:speed[-1]].copy())
pose_3d_cropped.append(pose_3d[start:end:speed[-1]].copy())
next_speed = min(
max(
speed[-1] + np.random.randint(
low=self.acc_range[0], high=self.acc_range[1] + 1),
1), self.speed_range[1])
speed.append(next_speed)
start = end
speed = speed[:-1] # The last speed not used
else: # Constant video speed
bps = []
speed = np.random.randint(low=self.speed_range[0],
high=self.speed_range[1] + 1)
pose_2d_cropped.append(
pose_2d[start:start + self.max_seq_len * speed:speed].copy())
pose_3d_cropped.append(
pose_3d[start:start + self.max_seq_len * speed:speed].copy())
pose_2d_cropped = np.concatenate(pose_2d_cropped, 0)
pose_3d_cropped = np.concatenate(pose_3d_cropped, 0)
# Pack data
data = {
'pose_2d': pose_2d_cropped,
'pose_3d': pose_3d_cropped,
'camera_view': camera_view,
'speed': speed,
'breakpoints': bps
}
return data
def get_seq_len(self):
seq_len = np.random.randint(low=self.max_seq_len - 2 * self.delta_len,
high=self.max_seq_len + 1)
return seq_len
def remove_joints(self, joints_to_remove):
kept_joints = self.skeleton.remove_joints(joints_to_remove)
for seq_i in range(len(self.data['pose']['3d'])):
self.data['pose']['3d'][seq_i] = self.data['pose']['3d'][
seq_i][:, kept_joints]
self.data['pose']['2d'][seq_i] = self.data['pose']['2d'][
seq_i][:, kept_joints]
def __len__(self):
return len(self.data['pose']['2d'])
