T2M-GPT 코드 분석 (2) HumanML3D preprocessing 정복하기

HumanML3D preprocessing

T2M-GPT를 보면 HumanML3D Dataset을 사용하는 것을 알 수가 있는데 여기에서 Extract and Process Data를 motion_representation.ipynb를 확인하고자 한다.
해당 소스 코드를 봤을 때 처음에는 그냥 3D Pose를 6D로 변경시켜서 concat한 다음 data를 만들었다고 생각했으나 완전히 틀린 생각이였다.

New Ground Truth 3D Pose

def uniform_skeleton(positions, target_offset):
    src_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')
    src_offset = src_skel.get_offsets_joints(torch.from_numpy(positions[0]))
    src_offset = src_offset.numpy()
    tgt_offset = target_offset.numpy()
    # print(src_offset)
    # print(tgt_offset)
    '''Calculate Scale Ratio as the ratio of legs'''
    src_leg_len = np.abs(src_offset[l_idx1]).max() + np.abs(src_offset[l_idx2]).max()
    tgt_leg_len = np.abs(tgt_offset[l_idx1]).max() + np.abs(tgt_offset[l_idx2]).max()

    scale_rt = tgt_leg_len / src_leg_len
    # print(scale_rt)
    src_root_pos = positions[:, 0]
    tgt_root_pos = src_root_pos * scale_rt

    '''Inverse Kinematics'''
    quat_params = src_skel.inverse_kinematics_np(positions, face_joint_indx)
    # print(quat_params.shape)

    '''Forward Kinematics'''
    src_skel.set_offset(target_offset)
    new_joints = src_skel.forward_kinematics_np(quat_params, tgt_root_pos)
    return new_joints

해당 소스코드는 기존의 SMPL로 만든 3D Pose를 tgt_offset target offset을 통해 IK+FK기법을 통해 target에 맞는 pose로 바꾸는 역할을 한다.

def process_file을 보면 Uniform Skeleton을 거친 후에 Put on Floor를 거치는데

    floor_height = positions.min(axis=0).min(axis=0)[1]
    positions[:, :, 1] -= floor_height

가장 작은 값을 찾아서 positions의 y값들을 그만큼 빼주는 작업을 해 바닥을 거치도록 만든다.
이코드를 분석해 봤을 때 모든 motion에는 바닥을 거친다는 가정을 통해 만든 것 같다.

XZ at origin의 코드를 분석해보면 첫번째 motion에 대한 pose를 가져와 거기서 pelvis에 해당하는 root pose를 가져와 x,z를 모든 motion에 뺀다.

    root_pos_init = positions[0]
    root_pose_init_xz = root_pos_init[0] * np.array([1, 0, 1])
    positions = positions - root_pose_init_xz

이 이유는 3d pose & mesh estimation에서 보면 알 듯 evaluation을 진행할 때 root pose를 항상 빼준 후에 metric을 진행한다. motion 또한 항상 처음 motion의 위치를 고정해주고자 만든 code인 것 같다.

All initially face Z+의 코드를 분석해보면 r_hip, l_hip, sdr_r, sdr_l을 가져와 서로 hip, sdr(어깨)를 빼준 후에 더해준다. 그 후에 x,y,z에 해당하는 값들을 제곱, 루트를 진행해 각 길이에 나눠준다. 그 후에 [0,1,0]을 외적을 구하면 사람이 진행하는 방향에 대한 forward_init이 나오게 된다. forward init을 구한 코드는 다음과 같다.

    r_hip, l_hip, sdr_r, sdr_l = face_joint_indx
    across1 = root_pos_init[r_hip] - root_pos_init[l_hip]
    across2 = root_pos_init[sdr_r] - root_pos_init[sdr_l]
    across = across1 + across2
    across = across / np.sqrt((across ** 2).sum(axis=-1))[..., np.newaxis]
    
    # forward (3,), rotate around y-axis
    forward_init = np.cross(np.array([[0, 1, 0]]), across, axis=-1)
    # forward (3,)
    forward_init = forward_init / np.sqrt((forward_init ** 2).sum(axis=-1))[..., np.newaxis]

이 preprocessing을 봤을 때 y가 floor, forward init을 구하기 위해 사용한 axis라고 가정하는 것을 봤을 때 y값이 하늘 x,z값이 걸어다니는 axis라고 예측할 수 있었다.

그 다음 아래 코드와 같이 root_quat_init이라고 구하게 되는데

    target = np.array([[0, 0, 1]])
    root_quat_init = qbetween_np(forward_init, target)
    root_quat_init = np.ones(positions.shape[:-1] + (4,)) * root_quat_init

qbetwwen_np의 코드는

def qbetween(v0, v1):
    '''
    find the quaternion used to rotate v0 to v1
    '''
    assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)'
    assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)'

    v = torch.cross(v0, v1)
    w = torch.sqrt((v0 ** 2).sum(dim=-1, keepdim=True) * (v1 ** 2).sum(dim=-1, keepdim=True)) + (v0 * v1).sum(dim=-1,
                                                                                                              keepdim=True)
    return qnormalize(torch.cat([w, v], dim=-1))


def qbetween_np(v0, v1):
    '''
    find the quaternion used to rotate v0 to v1
    '''
    assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)'
    assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)'

    v0 = torch.from_numpy(v0).float()
    v1 = torch.from_numpy(v1).float()
    return qbetween(v0, v1).numpy()

qbetween, qbetween_np로 구성되어 있는데 벡터 v0, v1의 외적 및 스칼라 값을 나타내는 4차원으로 구성된다고 생각하면 될 것 같다.
따라서 코드의 의미는 z값과 스칼라 값 + forward_init의 외적이라고 생각하면 될 것 같다.

이제 ground truth의 positions을 구하게 되는데 qrot_np를 root_quat_init, positions를 통해 구하게 된다.

    positions_b = positions.copy()

    positions = qrot_np(root_quat_init, positions)

    #     plot_3d_motion("./positions_2.mp4", kinematic_chain, positions, 'title', fps=20)

    '''New ground truth positions'''
    global_positions = positions.copy()

여기서 qrot_np는 다음과 같은 코드를 거치게 되는데

def qrot(q, v):
    """
    Rotate vector(s) v about the rotation described by quaternion(s) q.
    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
    where * denotes any number of dimensions.
    Returns a tensor of shape (*, 3).
    """
    assert q.shape[-1] == 4
    assert v.shape[-1] == 3
    assert q.shape[:-1] == v.shape[:-1]

    original_shape = list(v.shape)
    # print(q.shape)
    q = q.contiguous().view(-1, 4)
    v = v.contiguous().view(-1, 3)

    qvec = q[:, 1:]
    uv = torch.cross(qvec, v, dim=1)
    uuv = torch.cross(qvec, uv, dim=1)
    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)

    def qrot_np(q, v):
    #print(q.shape,v.shape)
    q = torch.from_numpy(q).contiguous().float()
    v = torch.from_numpy(v).contiguous().float()
    #print(q.shape,v.shape)
    return qrot(q, v).numpy()

아까 root_quat_init는 스칼라 값 + 외적을 구했는데 거기서 외적만 가지고 와서 다시 postion과 외적을 거친다. 이때 구한 rotation은 quaternion을 통해 구한 roation이다.

Motion Representation preprocessing

ground truth를 새롭게 바꿧고 이제 motion representation을 거친다.
가장 먼저 def foot_detect를 거치는데

    def foot_detect(positions, thres):
        velfactor, heightfactor = np.array([thres, thres]), np.array([3.0, 2.0])

        feet_l_x = (positions[1:, fid_l, 0] - positions[:-1, fid_l, 0]) ** 2
        feet_l_y = (positions[1:, fid_l, 1] - positions[:-1, fid_l, 1]) ** 2
        feet_l_z = (positions[1:, fid_l, 2] - positions[:-1, fid_l, 2]) ** 2
        #     feet_l_h = positions[:-1,fid_l,1]
        #     feet_l = (((feet_l_x + feet_l_y + feet_l_z) < velfactor) & (feet_l_h < heightfactor)).astype(np.float)
        feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float32)

        feet_r_x = (positions[1:, fid_r, 0] - positions[:-1, fid_r, 0]) ** 2
        feet_r_y = (positions[1:, fid_r, 1] - positions[:-1, fid_r, 1]) ** 2
        feet_r_z = (positions[1:, fid_r, 2] - positions[:-1, fid_r, 2]) ** 2
        #     feet_r_h = positions[:-1,fid_r,1]
        #     feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor) & (feet_r_h < heightfactor)).astype(np.float)
        feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor)).astype(np.float32)
        return feet_l, feet_r
    #
    feet_l, feet_r = foot_detect(positions, feet_thre)

feet_l_x = (positions[1:, fid_l, 0] - positions[:-1, fid_l, 0]) ** 2이 코드만 봤을 때 motion에서 첫번째를 제외한 나머지 postion들의 발 위치 - 이전 motion들의 발들의 x**2 값이라고 생각하면 될 것 같다.
feet_l은 0.002라고 정해논 thres값보다 크게되면 False로 두게 만드는 코드이다. threshold는 제곱승이기 때문에 좀 더 큰 값이라고 판단하면 될 것 같다. 이렇게 왼쪽 발 오른쪽 발에 대한 true/false 값을 만든다고 판단하면 될 것 같다.

그 다음으로 get_cont6d_params를 거치는데 아래의 코드와 같다.

    def get_cont6d_params(positions):
        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
        # (seq_len, joints_num, 4)
        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=True)

        '''Quaternion to continuous 6D'''
        cont_6d_params = quaternion_to_cont6d_np(quat_params)
        # (seq_len, 4)
        r_rot = quat_params[:, 0].copy()
        #     print(r_rot[0])
        '''Root Linear Velocity'''
        # (seq_len - 1, 3)
        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
        #     print(r_rot.shape, velocity.shape)
        velocity = qrot_np(r_rot[1:], velocity)
        '''Root Angular Velocity'''
        # (seq_len - 1, 4)
        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
        # (seq_len, joints_num, 4)
        return cont_6d_params, r_velocity, velocity, r_rot

skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")에서 n_raw_offsets은 paramUtil/t2m_raw_offsets을 의미하고 kinematic_chain은 paramUtil/t2m_kinematic_chain를 의미한다.
여기서 inverse_kinematics_np를 진행하는데 이는 common/skeleton/Skeleton/inverse_kinematics_np에 존재한다.

    def inverse_kinematics_np(self, joints, face_joint_idx, smooth_forward=False):
        assert len(face_joint_idx) == 4
        '''Get Forward Direction'''
        l_hip, r_hip, sdr_r, sdr_l = face_joint_idx
        across1 = joints[:, r_hip] - joints[:, l_hip]
        across2 = joints[:, sdr_r] - joints[:, sdr_l]
        across = across1 + across2
        across = across / np.sqrt((across**2).sum(axis=-1))[:, np.newaxis]
        # print(across1.shape, across2.shape)

        # forward (batch_size, 3)
        forward = np.cross(np.array([[0, 1, 0]]), across, axis=-1)
        if smooth_forward:
            forward = filters.gaussian_filter1d(forward, 20, axis=0, mode='nearest')
            # forward (batch_size, 3)
        forward = forward / np.sqrt((forward**2).sum(axis=-1))[..., np.newaxis]

        '''Get Root Rotation'''
        target = np.array([[0,0,1]]).repeat(len(forward), axis=0)
        root_quat = qbetween_np(forward, target)

        '''Inverse Kinematics'''
        # quat_params (batch_size, joints_num, 4)
        # print(joints.shape[:-1])
        quat_params = np.zeros(joints.shape[:-1] + (4,))
        # print(quat_params.shape)
        root_quat[0] = np.array([[1.0, 0.0, 0.0, 0.0]])
        quat_params[:, 0] = root_quat
        # quat_params[0, 0] = np.array([[1.0, 0.0, 0.0, 0.0]])
        for chain in self._kinematic_tree:
            R = root_quat
            for j in range(len(chain) - 1):
                # (batch, 3)
                u = self._raw_offset_np[chain[j+1]][np.newaxis,...].repeat(len(joints), axis=0)
                # print(u.shape)
                # (batch, 3)
                v = joints[:, chain[j+1]] - joints[:, chain[j]]
                v = v / np.sqrt((v**2).sum(axis=-1))[:, np.newaxis]
                # print(u.shape, v.shape)
                rot_u_v = qbetween_np(u, v)

                R_loc = qmul_np(qinv_np(R), rot_u_v)

                quat_params[:,chain[j + 1], :] = R_loc
                R = qmul_np(R, R_loc)

        return quat_params

코드만 봤을 때 위에서 진행한 quaternion 공식들을 이용한 것을 그대로 진행한 것을 볼 수 있다. 여기서 smooth_forward라고 있는데 gaussian_filter를 이용해 smooth하게 만들 것인지 안만들 것인지에 대한 판단이다.
하지만 여기서 get Root Rotation이라고 있는데 이점이 조금 다르게 z값에 대한 qbetween_np를 구한다.
그 후에 inverse kinematics를 진행한다. 오늘은 여기까지 정리하고 내일 마저 정리할 예정이다!