def positions_to_local_rot()

in src/definitions.py [0:0]

• 56 lines of code
• 12 McCabe index (conditional complexity)

    def positions_to_local_rot(self, p_g):
        r_l = torch.zeros_like(p_g)
        quat_g = torch.zeros((p_g.shape[0], p_g.shape[1], 4), device=p_g.device)
        quat_l = torch.zeros((p_g.shape[0], p_g.shape[1], 4), device=p_g.device)
        quat_g[:, :, 0] += 1.0
        quat_l[:, :, 0] += 1.0

        bone_offsets = torch.from_numpy(self.bone_offset_vector() / 100.0).float().to(p_g.device)
        child_idx_vector = self.child_idx_vector()
        parent_idx_vector = self.parent_idx_vector()

        depend_on_parent_idx = [1, 2, 3, 4]

        for joint_idx, parent_idx in enumerate(parent_idx_vector):
            parent_g = quat_g[:, parent_idx]

            use_par = joint_idx in depend_on_parent_idx

            my_pos = [torch.zeros_like(p_g[:, joint_idx])]
            ref_pos = [torch.zeros_like(p_g[:, joint_idx])]
            for child_idx in child_idx_vector[joint_idx]:
                ref_bone_dir = torch_tile(bone_offsets[child_idx][None, :].clone(), dim=0, n_tile=len(p_g))
                if torch.norm(ref_bone_dir) == 0.0:
                    continue
                my_pos.append(p_g[:, child_idx] - p_g[:, joint_idx])
                if use_par:
                    ref_bone_dir = quaternion.qrot(parent_g, ref_bone_dir)
                ref_pos.append(ref_bone_dir)

            finger_knuckles = []
            # extrimities_start = [Id.rupperarm, Id.lupperarm, Id.rupperleg, Id.lupperleg]
            extrimities_start = []
            if len(my_pos) == 2 and (joint_idx in extrimities_start or joint_idx in finger_knuckles):
                child_idx = child_idx_vector[joint_idx][0]
                grand_child_idx = child_idx_vector[child_idx][0]
                my_pos.append(p_g[:, grand_child_idx] - p_g[:, joint_idx])
                b_off = bone_offsets[grand_child_idx] + bone_offsets[child_idx]
                ref_bone_dir = torch_tile(b_off[None, :], dim=0, n_tile=len(p_g))
                ref_pos.append(ref_bone_dir)

            if len(my_pos) > 2:
                if use_par:
                    quat_l[:, joint_idx] = rotation_between_triangles(ref_pos[:3], my_pos[:3])
                    r_l[:, joint_idx] = quaternion.qeuler(quat_l[:, joint_idx])
                    quat_g[:, joint_idx] = quaternion.qmul(parent_g, quat_l[:, joint_idx])
                else:
                    quat_g[:, joint_idx] = rotation_between_triangles(ref_pos[:3], my_pos[:3])
                    quat_l[:, joint_idx] = quaternion.qmul(quaternion.inverse(parent_g), quat_g[:, joint_idx])
                    r_l[:, joint_idx] = quaternion.qeuler(quat_l[:, joint_idx])
            elif len(my_pos) == 2:
                ref_bone_dir = normalize_batch(ref_pos[1])
                child_idx = child_idx_vector[joint_idx][0]
                anim_bone_dir = normalize_batch(p_g[:, child_idx] - p_g[:, joint_idx])
                if use_par:
                    quat_l[:, joint_idx] = quaternion.quaternion_between_vectors_torch(ref_bone_dir, anim_bone_dir)
                    r_l[:, joint_idx] = quaternion.qeuler(quat_l[:, joint_idx])
                    quat_g[:, joint_idx] = quaternion.qmul(parent_g, quat_l[:, joint_idx])
                else:
                    quat_g[:, joint_idx] = quaternion.quaternion_between_vectors_torch(ref_bone_dir, anim_bone_dir)
                    quat_l[:, joint_idx] = quaternion.qmul(quaternion.inverse(parent_g), quat_g[:, joint_idx])
                    r_l[:, joint_idx] = quaternion.qeuler(quat_l[:, joint_idx])
            else:
                quat_g[:, joint_idx] = parent_g.clone()

        return rad2deg(r_l), quat_l