def encode(self,
               keypoints: np.ndarray,
               keypoints_visible: Optional[np.ndarray] = None) -> dict:
        """Encode keypoints into root heatmaps and keypoint displacement
        fields. Note that the original keypoint coordinates should be in the
        input image space.

        Args:
            keypoints (np.ndarray): Keypoint coordinates in shape (N, K, D)
            keypoints_visible (np.ndarray): Keypoint visibilities in shape
                (N, K)

        Returns:
            dict:
            - heatmaps (np.ndarray): The generated heatmap in shape
                (1, H, W) where [W, H] is the `heatmap_size`. If keypoint
                heatmaps are generated together, the shape is (K+1, H, W)
            - heatmap_weights (np.ndarray): The pixel-wise weight for heatmaps
                 which has same shape with `heatmaps`
            - displacements (np.ndarray): The generated displacement fields in
                shape (K*D, H, W). The vector on each pixels represents the
                displacement of keypoints belong to the associated instance
                from this pixel.
            - displacement_weights (np.ndarray): The pixel-wise weight for
                displacements which has same shape with `displacements`
        """

        if keypoints_visible is None:
            keypoints_visible = np.ones(keypoints.shape[:2], dtype=np.float32)

        # keypoint coordinates in heatmap
        _keypoints = keypoints / self.scale_factor

        # compute the root and scale of each instance
        roots, roots_visible = get_instance_root(_keypoints, keypoints_visible,
                                                 self.root_type)
        diagonal_lengths = get_diagonal_lengths(_keypoints, keypoints_visible)

        # discard the small instances
        roots_visible[diagonal_lengths < self.minimal_diagonal_length] = 0

        # generate heatmaps
        heatmaps, _ = generate_gaussian_heatmaps(
            heatmap_size=self.heatmap_size,
            keypoints=roots[:, None],
            keypoints_visible=roots_visible[:, None],
            sigma=self.sigma[0])
        heatmap_weights = self._get_heatmap_weights(
            heatmaps, bg_weight=self.background_weight)

        if self.generate_keypoint_heatmaps:
            keypoint_heatmaps, _ = generate_gaussian_heatmaps(
                heatmap_size=self.heatmap_size,
                keypoints=_keypoints,
                keypoints_visible=keypoints_visible,
                sigma=self.sigma[1])

            keypoint_heatmaps_weights = self._get_heatmap_weights(
                keypoint_heatmaps, bg_weight=self.background_weight)

            heatmaps = np.concatenate((keypoint_heatmaps, heatmaps), axis=0)
            heatmap_weights = np.concatenate(
                (keypoint_heatmaps_weights, heatmap_weights), axis=0)

        # generate displacements
        displacements, displacement_weights = \
            generate_displacement_heatmap(
                self.heatmap_size,
                _keypoints,
                keypoints_visible,
                roots,
                roots_visible,
                diagonal_lengths,
                self.sigma[0],
            )

        encoded = dict(
            heatmaps=heatmaps,
            heatmap_weights=heatmap_weights,
            displacements=displacements,
            displacement_weights=displacement_weights)

        return encoded
Encode keypoints into root heatmaps and keypoint displacement
        fields. Note that the original keypoint coordinates should be in the
        input image space.

        Args:
            keypoints (np.ndarray): Keypoint coordinates in shape (N, K, D)
            keypoints_visible (np.ndarray): Keypoint visibilities in shape
                (N, K)

        Returns:
            dict:
            - heatmaps (np.ndarray): The generated heatmap in shape
                (1, H, W) where [W, H] is the `heatmap_size`. If keypoint
                heatmaps are generated together, the shape is (K+1, H, W)
            - heatmap_weights (np.ndarray): The pixel-wise weight for heatmaps
                 which has same shape with `heatmaps`
            - displacements (np.ndarray): The generated displacement fields in
                shape (K*D, H, W). The vector on each pixels represents the
                displacement of keypoints belong to the associated instance
                from this pixel.
            - displacement_weights (np.ndarray): The pixel-wise weight for
                displacements which has same shape with `displacements`
        
mmpose/mmpose/codecs/spr.py