def batch_decode(self, batch_heatmaps: Tensor, batch_tags: Tensor
                     ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
        """Decode the keypoint coordinates from a batch of heatmaps and tagging
        heatmaps. The decoded keypoint coordinates are in the input image
        space.

        Args:
            batch_heatmaps (Tensor): Keypoint detection heatmaps in shape
                (B, K, H, W)
            batch_tags (Tensor): Tagging heatmaps in shape (B, C, H, W), where
                :math:`C=L*K`

        Returns:
            tuple:
            - batch_keypoints (List[np.ndarray]): Decoded keypoint coordinates
                of the batch, each is in shape (N, K, D)
            - batch_scores (List[np.ndarray]): Decoded keypoint scores of the
                batch, each is in shape (N, K). It usually represents the
                confidience of the keypoint prediction
        """
        B, _, H, W = batch_heatmaps.shape
        assert batch_tags.shape[0] == B and batch_tags.shape[2:4] == (H, W), (
            f'Mismatched shapes of heatmap ({batch_heatmaps.shape}) and '
            f'tagging map ({batch_tags.shape})')

        # Heatmap NMS
        batch_heatmaps = batch_heatmap_nms(batch_heatmaps,
                                           self.decode_nms_kernel)

        # Get top-k in each heatmap and and convert to numpy
        batch_topk_vals, batch_topk_tags, batch_topk_locs = to_numpy(
            self._get_batch_topk(
                batch_heatmaps, batch_tags, k=self.decode_topk))

        # Group keypoint candidates into groups (instances)
        batch_groups = self._group_keypoints(batch_topk_vals, batch_topk_tags,
                                             batch_topk_locs)

        # Convert to numpy
        batch_heatmaps_np = to_numpy(batch_heatmaps)
        batch_tags_np = to_numpy(batch_tags)

        # Refine the keypoint prediction
        batch_keypoints = []
        batch_keypoint_scores = []
        for i, (groups, heatmaps, tags) in enumerate(
                zip(batch_groups, batch_heatmaps_np, batch_tags_np)):

            keypoints, scores = groups[..., :-1], groups[..., -1]

            if keypoints.size > 0:
                # identify missing keypoints
                keypoints, scores = self._fill_missing_keypoints(
                    keypoints, scores, heatmaps, tags)

                # refine keypoint coordinates according to heatmap distribution
                if self.use_udp:
                    keypoints = refine_keypoints_dark_udp(
                        keypoints,
                        heatmaps,
                        blur_kernel_size=self.decode_gaussian_kernel)
                else:
                    keypoints = refine_keypoints(keypoints, heatmaps)

            batch_keypoints.append(keypoints)
            batch_keypoint_scores.append(scores)

        # restore keypoint scale
        batch_keypoints = [
            kpts * self.scale_factor for kpts in batch_keypoints
        ]

        return batch_keypoints, batch_keypoint_scores
Decode the keypoint coordinates from a batch of heatmaps and tagging
        heatmaps. The decoded keypoint coordinates are in the input image
        space.

        Args:
            batch_heatmaps (Tensor): Keypoint detection heatmaps in shape
                (B, K, H, W)
            batch_tags (Tensor): Tagging heatmaps in shape (B, C, H, W), where
                :math:`C=L*K`

        Returns:
            tuple:
            - batch_keypoints (List[np.ndarray]): Decoded keypoint coordinates
                of the batch, each is in shape (N, K, D)
            - batch_scores (List[np.ndarray]): Decoded keypoint scores of the
                batch, each is in shape (N, K). It usually represents the
                confidience of the keypoint prediction
        
mmpose/mmpose/codecs/associative_embedding.py