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()