def xyz2qrel(xyz_path, out_path, part): """Convert xyz data to quaternion data (angle between bone and Z axis) Args: xyz_path: (str) directory of processed xyz data out_path: (str) directory to save converted quaternion data part: (str) 'train' or 'eval' """ q_data_path = os.path.join(out_path, '{}_data.npy'.format(part)) q_label_path = os.path.join(out_path, '{}_label.pkl'.format(part)) xyz_data_path = os.path.join(xyz_path, '{}_data.npy'.format(part)) xyz_label_path = os.path.join(xyz_path, '{}_label.pkl'.format(part)) if not os.path.exists(q_label_path): copyfile(xyz_label_path, q_label_path) # Read xyz data (num_samples, 3, num_frame, num_joint, max_body) xyz_data = np.load(xyz_data_path, mmap_mode='r') num_sample, _, max_frame, num_joint, max_body = xyz_data.shape # Get num_frame with open(q_label_path, 'rb') as f: _, _, num_frame = pickle.load(f) # Open memory to write fp = open_memmap(q_data_path, dtype='float32', mode='w+', shape=(num_sample, 4, max_frame, num_joint, max_body)) fp[:, 0, ...] = 1. # Convert xyz to quaternion and save sk = Skeleton(ntu_info.PARENTS) xyz = xyz_data.transpose( 0, 2, 4, 1, 3) # (num_samples, num_frame, max_body, 3, num_joint) for i in tqdm(range(num_sample), ascii=True): fp[i][:, :num_frame[i], :, :] = sk.xyz2qrel( xyz[i][:num_frame[i]]).transpose( 2, 0, 3, 1) # (num_samples, 4, num_frame, num_joint, max_body)
def random_rotate_fromXYZ(xin): from data.fpha import fpha_info from utils.skel_utils.skeleton import Skeleton sk = Skeleton(fpha_info.PARENTS) xin = np.array(xin) xin = xin.transpose(0, 2, 3, 1) # (B,F,J,3) xin = sk.xyz_rotate(xin).transpose(0,1,3,2) # (B,F,3,J) xin = torch.from_numpy(sk.xyz2qrel(xin)).float().permute(0,2,1,3).unsqueeze(-1) # (B,4,F,J,1) return xin
def ntu_avg_bone_length(): from data.ntu.feeder import Feeder from data.ntu import ntu_info train_data_path = '/media/xuan/ssd/data/NTU-RGB-D/xyz/xsub/train_data.npy' train_label_path = '/media/xuan/ssd/data/NTU-RGB-D/xyz/xsub/train_label.pkl' train_data = Feeder(train_data_path, train_label_path, num_samples=-1, num_frames=300, mmap=True) xyz = np.zeros((3, 25)) total_frames = 0 for i in range(len(train_data)): seq, _ = train_data[i] for j in range(train_data.num_frames_data[i]): xyz += seq[:, j, :] total_frames += train_data.num_frames_data[i] xyz /= total_frames sk = Skeleton(parents=ntu_info.PARENTS) print(sk.compute_bone_lens(xyz))
def dhg_avg_bone_length(): from data.dhg.feeder import Feeder from data.dhg import dhg_info train_data_path = '/media/xuan/ssd/data/dhg_processed/xyz/train_data.npy' train_label_path = '/media/xuan/ssd/data/dhg_processed/xyz/train_14_label.pkl' train_data = Feeder(train_data_path, train_label_path, num_samples=-1, num_frames=200, mmap=True) xyz = np.zeros((3, 22)) total_frames = 0 for i in range(len(train_data)): seq, _ = train_data[i] for j in range(train_data.num_frames_data[i]): xyz += seq[:, j, :] total_frames += train_data.num_frames_data[i] xyz /= total_frames sk = Skeleton(parents=dhg_info.PARENTS) print(sk.compute_bone_lens(xyz))
def ntu_anim_qabs(seed=None): """Create an animation of skeleton in ntu dataset """ root = '/media/xuan/ssd/data/NTU-RGB-D-pre' benchmark = 'xsub' modality = 'quaternion' data_path = os.path.join(root, modality, benchmark, 'train_data.npy') label_path = os.path.join(root, modality, benchmark, 'train_label.pkl') data = Feeder(data_path, label_path, num_samples=-1, num_frames=20, mmap=True) sk = Skeleton(parents=ntu_info.PARENTS) # Prepare data np.random.seed(seed) index = np.random.randint(len(data)) quaternion, label = data[index] quaternion = quaternion[..., 0] xyz = sk.qabs2xyz(quaternion.transpose(1, 0, 2), ntu_info.AVG_BONE_LENS).transpose(1, 0, 2) # Prepare fig and axe fig = plt.figure(ntu_info.LABEL_NAMES[label]) ax = fig.add_subplot(1, 1, 1, projection='3d', xlabel='X', ylabel='Y', zlabel='Z') ax.view_init(azim=90, elev=-70) # Add animation anim = skeleton_anim(fig, ax, xyz, ntu_info.PARENTS, fps=10, colors=COLORS) anim.save('data/plots/ntu/qabs.gif', writer='imagemagick', fps=10) # Show() plt.show()
import torch.nn as nn import torch.optim as optim import yaml import argparse import importlib from shutil import copyfile from time import time import datetime import numpy as np from tqdm import tqdm from utils.skel_utils.skeleton import Skeleton from utils.quaternion import q_angle from data.ntu import ntu_info sk = Skeleton(ntu_info.PARENTS) def worker_init_fn(worker_id): np.random.seed(1346) def random_rotate(xin,y_only=False): if type(xin)==list: xin = [x.permute(0,4,2,3,1) for x in xin] s = xin[0].shape[0] for i in range(s): xin[0][i,...,1:]=torch.from_numpy(sk.xyz_rotate(xin[0][i,...,1:],y_only=y_only)).float() xin[1][i,...,1:]=torch.from_numpy(sk.xyz_rotate(xin[1][i,...,1:],new_rotate=False)).float() return [x.permute(0,4,2,3,1) for x in xin] else: xin = xin.permute(0,4,2,3,1) # B, M, F, J, 4 s = xin.shape[0]
def fpha_compare(): """Create side by side animation of xyz and quaternion reconstruction """ root = '/media/xuan/ssd/data/fpha_processed' num_frames = 20 xyz_data_path = os.path.join(root, 'xyz', 'train_data.npy') xyz_label_path = os.path.join(root, 'xyz', 'train_label.pkl') xyz_data = Feeder(xyz_data_path, xyz_label_path, num_samples=-1, num_frames=num_frames, mmap=True) q_data_path = os.path.join(root, 'quaternion', 'train_data.npy') q_label_path = os.path.join(root, 'quaternion', 'train_label.pkl') q_data = Feeder(q_data_path, q_label_path, num_samples=-1, num_frames=num_frames, mmap=True) qrec_data_path = os.path.join(root, 'qrec', 'train_data.npy') qrec_label_path = os.path.join(root, 'qrec', 'train_label.pkl') qrec_data = Feeder(qrec_data_path, qrec_label_path, num_samples=-1, num_frames=num_frames, mmap=True) colors = [''] * (NUM_JOINTS - 1) for i in range(NUM_JOINTS - 1): if i + 1 in fpha_info.T: colors[i] = 'r' elif i + 1 in fpha_info.I: colors[i] = 'g' elif i + 1 in fpha_info.M: colors[i] = 'b' elif i + 1 in fpha_info.R: colors[i] = 'c' elif i + 1 in fpha_info.P: colors[i] = 'm' np.random.seed() sk = Skeleton(parents=fpha_info.PARENTS) for _ in range(20): # Prepare data i = np.random.randint(len(xyz_data)) # xyz xyz, label = xyz_data[i] # Loaded qrec qrec, _ = qrec_data[i] # Loaded qrec # Computed from loaded quaternion q, _ = q_data[i] qrec_computed = sk.qu2xyz(q.transpose(1, 0, 2), fpha_info.AVG_BONE_LENS) qrec_computed = qrec_computed.transpose(1, 0, 2) # Prepare fig and axes title = fpha_info.LABEL_NAMES[label] + ' (xyz)' fig1 = plt.figure(title, figsize=(5, 5)) ax1 = fig1.add_subplot(1, 1, 1, projection='3d', xlabel='X', ylabel='Y', zlabel='Z') ax1.set_title(title, loc='left') ax1.view_init(azim=-30, elev=-60) title = fpha_info.LABEL_NAMES[label] + ' (qrec loaded)' fig2 = plt.figure(title, figsize=(5, 5)) ax2 = fig2.add_subplot(1, 1, 1, projection='3d', xlabel='X', ylabel='Y', zlabel='Z') ax2.set_title(title, loc='left') ax2.view_init(azim=-30, elev=-60) title = fpha_info.LABEL_NAMES[label] + ' (qrec from q)' fig3 = plt.figure(title, figsize=(5, 5)) ax3 = fig3.add_subplot(1, 1, 1, projection='3d', xlabel='X', ylabel='Y', zlabel='Z') ax3.set_title(title, loc='left') ax3.view_init(azim=-30, elev=-60) # Add animiation anim1 = skeleton_anim(fig1, ax1, xyz, fpha_info.PARENTS, fps=10, colors=colors) anim2 = skeleton_anim(fig2, ax2, qrec, fpha_info.PARENTS, fps=10, colors=colors) anim3 = skeleton_anim(fig3, ax3, qrec_computed, fpha_info.PARENTS, fps=10, colors=colors) anim1.save('data/plots/fpha/fpha_xyz.gif', writer='imagemagick', fps=10) anim2.save('data/plots/fpha/fpha_qrec_loaded.gif', writer='imagemagick', fps=10) anim3.save('data/plots/fpha/fpha_qrec_from_q.gif', writer='imagemagick', fps=10) # Show plt.show()