def _fill_missing_keypoints(self, keypoints: np.ndarray,
                                keypoint_scores: np.ndarray,
                                heatmaps: np.ndarray, tags: np.ndarray):
        """Fill the missing keypoints in the initial predictions.

        Args:
            keypoints (np.ndarray): Keypoint predictions in shape (N, K, D)
            keypoint_scores (np.ndarray): Keypint score predictions in shape
                (N, K), in which 0 means the corresponding keypoint is
                missing in the initial prediction
            heatmaps (np.ndarry): Heatmaps in shape (K, H, W)
            tags (np.ndarray): Tagging heatmaps in shape (C, H, W) where
                C=L*K

        Returns:
            tuple:
            - keypoints (np.ndarray): Keypoint predictions with missing
                ones filled
            - keypoint_scores (np.ndarray): Keypoint score predictions with
                missing ones filled
        """

        N, K = keypoints.shape[:2]
        H, W = heatmaps.shape[1:]
        L = tags.shape[0] // K
        keypoint_tags = [tags[k::K] for k in range(K)]

        for n in range(N):
            # Calculate the instance tag (mean tag of detected keypoints)
            _tag = []
            for k in range(K):
                if keypoint_scores[n, k] > 0:
                    x, y = keypoints[n, k, :2].astype(np.int64)
                    x = np.clip(x, 0, W - 1)
                    y = np.clip(y, 0, H - 1)
                    _tag.append(keypoint_tags[k][:, y, x])

            tag = np.mean(_tag, axis=0)
            tag = tag.reshape(L, 1, 1)
            # Search maximum response of the missing keypoints
            for k in range(K):
                if keypoint_scores[n, k] > 0:
                    continue
                dist_map = np.linalg.norm(
                    keypoint_tags[k] - tag, ord=2, axis=0)
                cost_map = np.round(dist_map) * 100 - heatmaps[k]  # H, W
                y, x = np.unravel_index(np.argmin(cost_map), shape=(H, W))
                keypoints[n, k] = [x, y]
                keypoint_scores[n, k] = heatmaps[k, y, x]

        return keypoints, keypoint_scores
Fill the missing keypoints in the initial predictions.

        Args:
            keypoints (np.ndarray): Keypoint predictions in shape (N, K, D)
            keypoint_scores (np.ndarray): Keypint score predictions in shape
                (N, K), in which 0 means the corresponding keypoint is
                missing in the initial prediction
            heatmaps (np.ndarry): Heatmaps in shape (K, H, W)
            tags (np.ndarray): Tagging heatmaps in shape (C, H, W) where
                C=L*K

        Returns:
            tuple:
            - keypoints (np.ndarray): Keypoint predictions with missing
                ones filled
            - keypoint_scores (np.ndarray): Keypoint score predictions with
                missing ones filled
        
mmpose/mmpose/codecs/associative_embedding.py