| 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 |
| 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) |
| else: |
| return np.array(range(target_len)) % ori_len |
| 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 |
