Пример #1
0
def test_flip_regression():
    coords = np.random.rand(3, 3)
    flip_pairs = [[1, 2]]
    root = coords[:1]
    coords_flipped = coords.copy()
    coords_flipped[1] = coords[2]
    coords_flipped[2] = coords[1]
    coords_flipped[..., 0] = 2 * root[..., 0] - coords_flipped[..., 0]

    # static mode
    res_static = fliplr_regression(coords,
                                   flip_pairs,
                                   center_mode='static',
                                   center_x=root[0, 0])
    assert_array_almost_equal(res_static, coords_flipped)

    # root mode
    res_root = fliplr_regression(coords,
                                 flip_pairs,
                                 center_mode='root',
                                 center_index=0)
    assert_array_almost_equal(res_root, coords_flipped)
Пример #2
0
    def inference_model(self, x, flip_pairs=None):
        """Inference function.

        Returns:
            output_regression (np.ndarray): Output regression.

        Args:
            x (torch.Tensor[N, K, 2]): Input features.
            flip_pairs (None | list[tuple()):
                Pairs of keypoints which are mirrored.
        """
        output = self.forward(x)

        if flip_pairs is not None:
            output_regression = fliplr_regression(
                output.detach().cpu().numpy(),
                flip_pairs,
                center_mode='static',
                center_x=0)
        else:
            output_regression = output.detach().cpu().numpy()
        return output_regression