File size: 3,867 Bytes
bbde80b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import os
import torch
import torch.nn.functional as F
import numpy as np
import copy
def crop_scale(motion, scale_range=[1, 1]):
'''
Motion: [(M), T, 17, 3].
Normalize to [-1, 1]
'''
result = copy.deepcopy(motion)
valid_coords = motion[motion[..., 2]!=0][:,:2]
if len(valid_coords) < 4:
return np.zeros(motion.shape)
xmin = min(valid_coords[:,0])
xmax = max(valid_coords[:,0])
ymin = min(valid_coords[:,1])
ymax = max(valid_coords[:,1])
ratio = np.random.uniform(low=scale_range[0], high=scale_range[1], size=1)[0]
scale = max(xmax-xmin, ymax-ymin) * ratio
if scale==0:
return np.zeros(motion.shape)
xs = (xmin+xmax-scale) / 2
ys = (ymin+ymax-scale) / 2
result[...,:2] = (motion[..., :2]- [xs,ys]) / scale
result[...,:2] = (result[..., :2] - 0.5) * 2
result = np.clip(result, -1, 1)
return result
def crop_scale_3d(motion, scale_range=[1, 1]):
'''
Motion: [T, 17, 3]. (x, y, z)
Normalize to [-1, 1]
Z is relative to the first frame's root.
'''
result = copy.deepcopy(motion)
result[:,:,2] = result[:,:,2] - result[0,0,2]
xmin = np.min(motion[...,0])
xmax = np.max(motion[...,0])
ymin = np.min(motion[...,1])
ymax = np.max(motion[...,1])
ratio = np.random.uniform(low=scale_range[0], high=scale_range[1], size=1)[0]
scale = max(xmax-xmin, ymax-ymin) / ratio
if scale==0:
return np.zeros(motion.shape)
xs = (xmin+xmax-scale) / 2
ys = (ymin+ymax-scale) / 2
result[...,:2] = (motion[..., :2]- [xs,ys]) / scale
result[...,2] = result[...,2] / scale
result = (result - 0.5) * 2
return result
def flip_data(data):
"""
horizontal flip
data: [N, F, 17, D] or [F, 17, D]. X (horizontal coordinate) is the first channel in D.
Return
result: same
"""
left_joints = [4, 5, 6, 11, 12, 13]
right_joints = [1, 2, 3, 14, 15, 16]
flipped_data = copy.deepcopy(data)
flipped_data[..., 0] *= -1 # flip x of all joints
flipped_data[..., left_joints+right_joints, :] = flipped_data[..., right_joints+left_joints, :]
return flipped_data
def resample(ori_len, target_len, replay=False, randomness=True):
if replay:
if ori_len > target_len:
st = np.random.randint(ori_len-target_len)
return range(st, st+target_len) # Random clipping from sequence
else:
return np.array(range(target_len)) % ori_len # Replay padding
else:
if randomness:
even = np.linspace(0, ori_len, num=target_len, endpoint=False)
if ori_len < target_len:
low = np.floor(even)
high = np.ceil(even)
sel = np.random.randint(2, size=even.shape)
result = np.sort(sel*low+(1-sel)*high)
else:
interval = even[1] - even[0]
result = np.random.random(even.shape)*interval + even
result = np.clip(result, a_min=0, a_max=ori_len-1).astype(np.uint32)
else:
result = np.linspace(0, ori_len, num=target_len, endpoint=False, dtype=int)
return result
def split_clips(vid_list, n_frames, data_stride):
result = []
n_clips = 0
st = 0
i = 0
saved = set()
while i<len(vid_list):
i += 1
if i-st == n_frames:
result.append(range(st,i))
saved.add(vid_list[i-1])
st = st + data_stride
n_clips += 1
if i==len(vid_list):
break
if vid_list[i]!=vid_list[i-1]:
if not (vid_list[i-1] in saved):
resampled = resample(i-st, n_frames) + st
result.append(resampled)
saved.add(vid_list[i-1])
st = i
return result |