| import numpy as np |
| import os |
| import random |
| import torch |
| import copy |
| import torch.nn as nn |
| from lib.utils.tools import read_pkl |
| from lib.utils.utils_data import flip_data, crop_scale_3d |
| |
| class Augmenter2D(object): |
| """ |
| Make 2D augmentations on the fly. PyTorch batch-processing GPU version. |
| """ |
| def __init__(self, args): |
| self.d2c_params = read_pkl(args.d2c_params_path) |
| self.noise = torch.load(args.noise_path) |
| self.mask_ratio = args.mask_ratio |
| self.mask_T_ratio = args.mask_T_ratio |
| self.num_Kframes = 27 |
| self.noise_std = 0.002 |
|
|
| def dis2conf(self, dis, a, b, m, s): |
| f = a/(dis+a)+b*dis |
| shift = torch.randn(*dis.shape)*s + m |
| |
| shift = shift.to(dis.device) |
| return f + shift |
| |
| def add_noise(self, motion_2d): |
| a, b, m, s = self.d2c_params["a"], self.d2c_params["b"], self.d2c_params["m"], self.d2c_params["s"] |
| if "uniform_range" in self.noise.keys(): |
| uniform_range = self.noise["uniform_range"] |
| else: |
| uniform_range = 0.06 |
| motion_2d = motion_2d[:,:,:,:2] |
| batch_size = motion_2d.shape[0] |
| num_frames = motion_2d.shape[1] |
| num_joints = motion_2d.shape[2] |
| mean = self.noise['mean'].float() |
| std = self.noise['std'].float() |
| weight = self.noise['weight'][:,None].float() |
| sel = torch.rand((batch_size, self.num_Kframes, num_joints, 1)) |
| gaussian_sample = (torch.randn(batch_size, self.num_Kframes, num_joints, 2) * std + mean) |
| uniform_sample = (torch.rand((batch_size, self.num_Kframes, num_joints, 2))-0.5) * uniform_range |
| noise_mean = 0 |
| delta_noise = torch.randn(num_frames, num_joints, 2) * self.noise_std + noise_mean |
| |
| mean = mean.to(motion_2d.device) |
| std = std.to(motion_2d.device) |
| weight = weight.to(motion_2d.device) |
| gaussian_sample = gaussian_sample.to(motion_2d.device) |
| uniform_sample = uniform_sample.to(motion_2d.device) |
| sel = sel.to(motion_2d.device) |
| delta_noise = delta_noise.to(motion_2d.device) |
| |
| delta = gaussian_sample*(sel<weight) + uniform_sample*(sel>=weight) |
| delta_expand = torch.nn.functional.interpolate(delta.unsqueeze(1), [num_frames, num_joints, 2], mode='trilinear', align_corners=True)[:,0] |
| delta_final = delta_expand + delta_noise |
| motion_2d = motion_2d + delta_final |
| dx = delta_final[:,:,:,0] |
| dy = delta_final[:,:,:,1] |
| dis2 = dx*dx+dy*dy |
| dis = torch.sqrt(dis2) |
| conf = self.dis2conf(dis, a, b, m, s).clip(0,1).reshape([batch_size, num_frames, num_joints, -1]) |
| return torch.cat((motion_2d, conf), dim=3) |
| |
| def add_mask(self, x): |
| ''' motion_2d: (N,T,17,3) |
| ''' |
| N,T,J,C = x.shape |
| mask = torch.rand(N,T,J,1, dtype=x.dtype, device=x.device) > self.mask_ratio |
| mask_T = torch.rand(1,T,1,1, dtype=x.dtype, device=x.device) > self.mask_T_ratio |
| x = x * mask * mask_T |
| return x |
| |
| def augment2D(self, motion_2d, mask=False, noise=False): |
| if noise: |
| motion_2d = self.add_noise(motion_2d) |
| if mask: |
| motion_2d = self.add_mask(motion_2d) |
| return motion_2d |
| |
| class Augmenter3D(object): |
| """ |
| Make 3D augmentations when dataloaders get items. NumPy single motion version. |
| """ |
| def __init__(self, args): |
| self.flip = args.flip |
| if hasattr(args, "scale_range_pretrain"): |
| self.scale_range_pretrain = args.scale_range_pretrain |
| else: |
| self.scale_range_pretrain = None |
| |
| def augment3D(self, motion_3d): |
| if self.scale_range_pretrain: |
| motion_3d = crop_scale_3d(motion_3d, self.scale_range_pretrain) |
| if self.flip and random.random()>0.5: |
| motion_3d = flip_data(motion_3d) |
| return motion_3d |
