def flip_heatmaps(heatmaps: Tensor,
                  flip_indices: Optional[List[int]] = None,
                  flip_mode: str = 'heatmap',
                  shift_heatmap: bool = True):
    """Flip heatmaps for test-time augmentation.

    Args:
        heatmaps (Tensor): The heatmaps to flip. Should be a tensor in shape
            [B, C, H, W]
        flip_indices (List[int]): The indices of each keypoint's symmetric
            keypoint. Defaults to ``None``
        flip_mode (str): Specify the flipping mode. Options are:

            - ``'heatmap'``: horizontally flip the heatmaps and swap heatmaps
                of symmetric keypoints according to ``flip_indices``
            - ``'udp_combined'``: similar to ``'heatmap'`` mode but further
                flip the x_offset values
            - ``'offset'``: horizontally flip the offset fields and swap
                heatmaps of symmetric keypoints according to
                ``flip_indices``. x_offset values are also reversed
        shift_heatmap (bool): Shift the flipped heatmaps to align with the
            original heatmaps and improve accuracy. Defaults to ``True``

    Returns:
        Tensor: flipped heatmaps in shape [B, C, H, W]
    """

    if flip_mode == 'heatmap':
        heatmaps = heatmaps.flip(-1)
        if flip_indices is not None:
            assert len(flip_indices) == heatmaps.shape[1]
            heatmaps = heatmaps[:, flip_indices]
    elif flip_mode == 'udp_combined':
        B, C, H, W = heatmaps.shape
        heatmaps = heatmaps.view(B, C // 3, 3, H, W)
        heatmaps = heatmaps.flip(-1)
        if flip_indices is not None:
            assert len(flip_indices) == C // 3
            heatmaps = heatmaps[:, flip_indices]
        heatmaps[:, :, 1] = -heatmaps[:, :, 1]
        heatmaps = heatmaps.view(B, C, H, W)

    elif flip_mode == 'offset':
        B, C, H, W = heatmaps.shape
        heatmaps = heatmaps.view(B, C // 2, -1, H, W)
        heatmaps = heatmaps.flip(-1)
        if flip_indices is not None:
            assert len(flip_indices) == C // 2
            heatmaps = heatmaps[:, flip_indices]
        heatmaps[:, :, 0] = -heatmaps[:, :, 0]
        heatmaps = heatmaps.view(B, C, H, W)

    else:
        raise ValueError(f'Invalid flip_mode value "{flip_mode}"')

    if shift_heatmap:
        # clone data to avoid unexpected in-place operation when using CPU
        heatmaps[..., 1:] = heatmaps[..., :-1].clone()

    return heatmaps
Flip heatmaps for test-time augmentation.

    Args:
        heatmaps (Tensor): The heatmaps to flip. Should be a tensor in shape
            [B, C, H, W]
        flip_indices (List[int]): The indices of each keypoint's symmetric
            keypoint. Defaults to ``None``
        flip_mode (str): Specify the flipping mode. Options are:

            - ``'heatmap'``: horizontally flip the heatmaps and swap heatmaps
                of symmetric keypoints according to ``flip_indices``
            - ``'udp_combined'``: similar to ``'heatmap'`` mode but further
                flip the x_offset values
            - ``'offset'``: horizontally flip the offset fields and swap
                heatmaps of symmetric keypoints according to
                ``flip_indices``. x_offset values are also reversed
        shift_heatmap (bool): Shift the flipped heatmaps to align with the
            original heatmaps and improve accuracy. Defaults to ``True``

    Returns:
        Tensor: flipped heatmaps in shape [B, C, H, W]
    
mmpose/mmpose/models/utils/tta.py