def _rotate_batch()

in long_term/pose_network_long_term.py [0:0]

• 19 lines of code
• 1 McCabe index (conditional complexity)

    def _rotate_batch(self, buffer_rot, buffer_pos):
        """
        Data augmentation: rotate each sequence in this batch by a random angle.
        """
        # We cover a 4*pi range to make sure that we visit the full quaternion range
        random_angles = np.random.uniform(-2*np.pi, 2*np.pi, size=buffer_rot.shape[0])
        rrot = np.zeros((random_angles.shape[0], 4), dtype='float32')
        rrot[:, 0] = np.cos(random_angles/2)
        rrot[:, 2] = np.sin(random_angles/2)
        
        # Rotate root joints
        rrot_ = np.tile(rrot.reshape(-1, 1, 4), (1, buffer_rot.shape[1], 1))
        buffer_rot[:, :, :4] = qmul_np(rrot_, buffer_rot[:, :, :4])
        
        # Rotate positions
        nj = self.skeleton.num_joints()
        positions = buffer_pos[:, :, :nj*3].reshape(buffer_pos.shape[0], buffer_pos.shape[1], -1, 3)
        rrot_ = np.tile(rrot.reshape(-1, 1, 1, 4), (1, positions.shape[1], positions.shape[2], 1))
        buffer_pos[:, :, :nj*3] = qrot_np(rrot_, positions).reshape(buffer_pos[:, :, :nj*3].shape)
        
        # Rotate controls
        extra_angles = np.zeros((buffer_rot.shape[0], buffer_rot.shape[1], 2, 3), dtype='float32')
        extra_angles[:, :, 0, [0, 2]] = buffer_rot[:, :, [-2, -1]]
        extra_angles[:, :, 1, [0, 2]] = buffer_rot[:, :, [-4, -3]]
        rrot_ = np.tile(rrot.reshape(-1, 1, 1, 4), (1, extra_angles.shape[1], extra_angles.shape[2], 1))
        extra_angles = qrot_np(rrot_, extra_angles)
        buffer_rot[:, :, [-2, -1]] = extra_angles[:, :, 0, [0, 2]]
        buffer_rot[:, :, [-4, -3]] = extra_angles[:, :, 1, [0, 2]]
        
        return buffer_rot, buffer_pos