Spaces:
Running
on
Zero
Running
on
Zero
| import time | |
| import torch | |
| import torch.nn.functional as F | |
| import numpy as np | |
| def get_gaussian_kernel_1d(kernel_size, sigma, device): | |
| x = torch.arange(kernel_size).float() - (kernel_size // 2) | |
| g = torch.exp(-((x ** 2) / (2 * sigma ** 2))) | |
| g /= g.sum() | |
| kernel_weight = g.view(1, 1, -1).to(device) | |
| return kernel_weight | |
| def gaussian_filter_1d(data, kernel_size=3, sigma=1.0, weight=None): | |
| kernel_weight = get_gaussian_kernel_1d(kernel_size, sigma, data.device) if weight is None else weight | |
| data = F.pad(data, (kernel_size // 2, kernel_size // 2), mode='replicate') | |
| return F.conv1d(data, kernel_weight) | |
| def exponential_smoothing(x, d_x, alpha=0.5): | |
| return d_x + alpha * (x - d_x) | |
| class OneEuroFilter: | |
| # param setting: | |
| # realtime v2m: min_cutoff=1.0, beta=1.5 | |
| # motionshop 2d keypoint: min_cutoff=1.7, beta=0.3 | |
| def __init__(self, min_cutoff=1.0, beta=0.0, sampling_rate=30, d_cutoff=1.0, device='cuda'): | |
| self.min_cutoff = min_cutoff | |
| self.beta = beta | |
| self.sampling_rate = sampling_rate | |
| self.x_prev = None | |
| self.dx_prev = None | |
| self.d_cutoff = d_cutoff | |
| self.pi = torch.tensor(torch.pi, device=device) | |
| def smoothing_factor(self, cutoff): | |
| r = 2 * self.pi * cutoff / self.sampling_rate | |
| return r/ (1 + r) | |
| def filter(self, x): | |
| if self.x_prev is None: | |
| self.x_prev = x | |
| self.dx_prev = torch.zeros_like(x) | |
| return x | |
| a_d = self.smoothing_factor(self.d_cutoff) | |
| # 计算当前的速度 | |
| dx = (x - self.x_prev) * self.sampling_rate | |
| dx_hat = exponential_smoothing(dx, self.dx_prev, a_d) | |
| cutoff = self.min_cutoff + self.beta * torch.abs(dx_hat) | |
| a = self.smoothing_factor(cutoff) | |
| x_hat = exponential_smoothing(x, self.x_prev, a) | |
| self.x_prev = x_hat | |
| self.dx_prev = dx_hat | |
| return x_hat | |
| class Filter(): | |
| filter_factory = { | |
| 'gaussian': get_gaussian_kernel_1d, | |
| } | |
| def __init__(self, target_data, filter_type, filter_args): | |
| self.target_data = target_data | |
| self.filter = self.filter_factory[filter_type] | |
| self.filter_args = filter_args | |
| def process(self, network_outputs): | |
| filter_data = [] | |
| for human in network_outputs: | |
| filter_data.append(human[self.target_data]) | |
| filter_data = torch.stack(filter_data, dim=0) | |
| filter_data = self.filter(filter_data, **self.filter_args) | |
| for i, human in enumerate(network_outputs): | |
| human[self.target_data] = filter_data[i] | |
| if __name__ == '__main__': | |
| import argparse | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| from rot6d import rotation_6d_to_axis_angle, axis_angle_to_rotation_6d | |
| from humans import get_smplx_joint_names | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--data_path', type=str) | |
| parser.add_argument('--save_path', type=str) | |
| parser.add_argument('--name', type=str) | |
| args = parser.parse_args() | |
| fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(10, 8)) | |
| data_types = ['rotvec']#, 'j3d'] | |
| observe_keypoints = ['pelvis', 'head', 'left_wrist', 'left_knee'] | |
| joint_names = get_smplx_joint_names() | |
| data = np.load(f'{args.data_path}/shape_{args.name}.npy') | |
| fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(10, 8)) | |
| for i in range(2): | |
| for j in range(2): | |
| x = data[:, i*4 + j*2] | |
| print(x.shape) | |
| axs[i, j].plot(x) | |
| axs[i, j].set_title(f'{4 * i + 2 * j}') | |
| axs[i, j].plot(np.load(f'{args.data_path}/dist_{args.name}.npy')) | |
| plt.tight_layout() | |
| plt.savefig(f'{args.save_path}/shape_{args.name}.jpg') | |
| # for data_type in data_types: | |
| # data = np.load(f'{args.data_path}/{data_type}_{args.name}.npy') | |
| # fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(10, 8)) | |
| # for i in range(2): | |
| # for j in range(2): | |
| # # todo: something wrong here | |
| # filter = OneEuroFilter(min_cutoff=1, beta=0.01, sampling_rate=30, device='cuda:0') | |
| # x = data[:, joint_names.index(observe_keypoints[i*2+j])] #(F, 3) | |
| # print(x.shape) | |
| # x = axis_angle_to_rotation_6d(torch.tensor(x, device='cuda:0')) | |
| # x_filtered = x.clone() | |
| # start = time.time() | |
| # for k in range(x.shape[0]): | |
| # x_filtered[k] = filter.filter(x[k]) | |
| # print(x_filtered.shape[0]/(time.time()-start)) | |
| # # x_filtered = x.clone() | |
| # # a = 0.5 | |
| # # for k in range(1, x.shape[0]): | |
| # # x_filtered[k] = (1 - a) * x_filtered[k-1] + a * x[k] | |
| # #theta = np.linalg.norm(x, axis=-1) | |
| # #x = x / theta[..., None] | |
| # # f, n = x.shape | |
| # # x_filtered = gaussian_filter_1d(x.permute(1, 0).view(n, 1, -1), 11, 11) | |
| # # x_filtered = x_filtered.view(n, -1).permute(1, 0) | |
| # x = rotation_6d_to_axis_angle(x).cpu().numpy() | |
| # x_filtered = rotation_6d_to_axis_angle(x_filtered).cpu().numpy() | |
| # axs[i, j].plot(x[..., 0]) | |
| # axs[i, j].plot(x[..., 1]) | |
| # axs[i, j].plot(x[..., 2]) | |
| # axs[i, j].plot(x_filtered[..., 0]) | |
| # axs[i, j].plot(x_filtered[..., 1]) | |
| # axs[i, j].plot(x_filtered[..., 2]) | |
| # #axs[i, j].plot(theta) | |
| # axs[i, j].set_title(f'{observe_keypoints[i*2 + j]}') | |
| # plt.tight_layout() | |
| # plt.savefig(f'{args.save_path}/{data_type}_{args.name}.jpg') | |