- Code for converting BVH files.

ego_pose/data_process/convert_clip.py

@@ -0,0 +1,77 @@
+import os
+import sys
+
+sys.path.append(os.getcwd())
+
+from utils import *
+from utils.transformation import quaternion_from_euler
+from mujoco_py import load_model_from_path
+from mocap.skeleton import Skeleton
+from mocap.pose import load_bvh_file, interpolated_traj
+import pickle
+import glob
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--render', action='store_true', default=True)
+parser.add_argument('--model-id', type=str, default='humanoid_1205_orig')
+parser.add_argument('--mocap-id', type=str, default='0213')
+parser.add_argument('--range', type=int, default=(5, 20))
+parser.add_argument('--skt-id', type=str, default='take_01')
+parser.add_argument('--mocap-fr', type=int, default=120)
+parser.add_argument('--dt', type=float, default=1/30)
+parser.add_argument('--offset-z', type=float, default=0.0)
+args = parser.parse_args()
+
+model_file = 'assets/mujoco_models/%s.xml' % args.model_id
+model = load_model_from_path(model_file)
+body_qposaddr = get_body_qposaddr(model)
+
+skt_bvh = os.path.expanduser('datasets/traj/%s_%s.bvh' % (args.mocap_id, args.skt_id))
+exclude_bones = {'Thumb', 'Index', 'Middle', 'Ring', 'Pinky', 'End', 'Toe'}
+spec_channels = {'LeftForeArm': ['Zrotation'], 'RightForeArm': ['Zrotation'],
+                 'LeftLeg': ['Xrotation'], 'RightLeg': ['Xrotation']}
+skeleton = Skeleton()
+skeleton.load_from_bvh(skt_bvh, exclude_bones, spec_channels)
+
+
+def get_qpos(pose, bone_addr):
+    qpos = np.zeros(model.nq)
+    for bone_name, ind2 in body_qposaddr.items():
+        ind1 = bone_addr[bone_name]
+        if ind1[0] == 0:
+            trans = pose[ind1[0]:ind1[0] + 3].copy()
+            angles = pose[ind1[0] + 3:ind1[1]].copy()
+            quat = quaternion_from_euler(angles[0], angles[1], angles[2], 'rxyz')
+            qpos[ind2[0]:ind2[0] + 3] = trans
+            qpos[ind2[0] + 3:ind2[1]] = quat
+        else:
+            qpos[ind2[0]:ind2[1]] = pose[ind1[0]:ind1[1]]
+    return qpos
+
+
+def get_poses(bvh_file):
+    poses, bone_addr = load_bvh_file(bvh_file, skeleton)
+    poses_samp = interpolated_traj(poses, args.dt, mocap_fr=args.mocap_fr)
+    qpos_traj = []
+    for i in range(poses_samp.shape[0]):
+        cur_pose = poses_samp[i, :]
+        cur_qpos = get_qpos(cur_pose, bone_addr)
+        qpos_traj.append(cur_qpos)
+    qpos_traj = np.vstack(qpos_traj)
+    qpos_traj[:, 2] += args.offset_z
+    return qpos_traj
+
+
+bvh_files = glob.glob(os.path.expanduser('datasets/traj/%s_*.bvh' % args.mocap_id))
+bvh_files.sort()
+if args.range is not None:
+    bvh_files = bvh_files[args.range[0]: args.range[1]]
+print(bvh_files)
+for file in bvh_files:
+    print('extracting trajectory from %s' % file)
+    qpos_traj = get_poses(file)
+    name = os.path.splitext(os.path.basename(file))[0]
+    bvh_dir = os.path.dirname(file)
+    traj_file = '%s/%s_traj.p' % (bvh_dir, name)
+    pickle.dump(qpos_traj, open(traj_file, 'wb'))

mocap/pose.py

@@ -0,0 +1,93 @@
+import numpy as np
+import math
+from bvh import Bvh
+from utils.transformation import quaternion_slerp, quaternion_from_euler, euler_from_quaternion
+
+
+def load_amc_file(fname, scale):
+
+    with open(fname) as f:
+        content = f.readlines()
+
+    bone_addr = dict()
+    poses = []
+    cur_pos = None
+    fr = 1
+    for line in content:
+        line_words = line.split()
+        cmd = line_words[0]
+        if cmd == str(fr):
+            if cur_pos:
+                poses.append(np.array(cur_pos))
+            cur_pos = []
+            fr += 1
+        elif cur_pos is not None:
+            start_ind = len(cur_pos)
+            if cmd == 'root':
+                cur_pos += [float(word)*scale for word in line_words[1:4]]
+                cur_pos += [math.radians(float(word)) for word in line_words[4:]]
+            elif cmd == 'lfoot' or cmd == 'rfoot':
+                cur_pos += reversed([math.radians(float(word)) for word in line_words[1:]])
+                if len(cur_pos) < 3:
+                    cur_pos.insert(-1, 0.0)
+            else:
+                cur_pos += reversed([math.radians(float(word)) for word in line_words[1:]])
+            if fr == 2:
+                end_ind = len(cur_pos)
+                bone_addr[cmd] = (start_ind, end_ind)
+
+    if cur_pos:
+        poses.append(np.array(cur_pos))
+    poses = np.vstack(poses)
+    return poses, bone_addr
+
+
+def load_bvh_file(fname, skeleton):
+    with open(fname) as f:
+        mocap = Bvh(f.read())
+
+    # build bone_addr
+    bone_addr = dict()
+    start_ind = 0
+    for bone in skeleton.bones:
+        end_ind = start_ind + len(bone.channels)
+        bone_addr[bone.name] = (start_ind, end_ind)
+        start_ind = end_ind
+    dof_num = start_ind
+
+    poses = np.zeros((mocap.nframes, dof_num))
+    for i in range(mocap.nframes):
+        for bone in skeleton.bones:
+            trans = np.array(mocap.frame_joint_channels(i, bone.name, bone.channels))
+            if bone == skeleton.root:
+                trans[:3] *= skeleton.len_scale
+                trans[3:6] = np.deg2rad(trans[3:6])
+            else:
+                trans = np.deg2rad(trans)
+            start_ind, end_ind = bone_addr[bone.name]
+            poses[i, start_ind:end_ind] = trans
+
+    return poses, bone_addr
+
+
+def lin_interp(pose1, pose2, t):
+    pose_t = (1 - t) * pose1 + t * pose2
+    return pose_t
+
+
+def interpolated_traj(poses, sample_t=0.030, mocap_fr=120):
+    N = poses.shape[0]
+    T = float(N-1)/mocap_fr
+    num = int(math.floor(T/sample_t))
+    sampling_times = np.arange(num+1)*sample_t*mocap_fr
+
+    poses_samp = []
+    for t in sampling_times:
+        start = int(math.floor(t))
+        end = min(int(math.ceil(t)), poses.shape[0] - 1)
+        poses_samp.append(lin_interp(poses[start, :], poses[end, :], t-math.floor(t)))
+    poses_samp = np.vstack(poses_samp)
+
+    return poses_samp
+
+