mmpose/datasets/transforms/common_transforms.py

# Copyright (c) OpenMMLab. All rights reserved.
import warnings
from copy import deepcopy
from typing import Dict, List, Optional, Sequence, Tuple, Union

import cv2
import mmcv
import mmengine
import numpy as np
from mmcv.image import imflip
from mmcv.transforms import BaseTransform
from mmcv.transforms.utils import avoid_cache_randomness, cache_randomness
from mmengine import is_list_of
from mmengine.dist import get_dist_info
from scipy.stats import truncnorm
from scipy.ndimage import distance_transform_edt

from mmpose.codecs import *  # noqa: F401, F403
from mmpose.registry import KEYPOINT_CODECS, TRANSFORMS
from mmpose.structures.bbox import bbox_xyxy2cs, flip_bbox, bbox_cs2xyxy
from mmpose.structures.keypoint import flip_keypoints
from mmpose.utils.typing import MultiConfig

from pycocotools import mask as Mask

try:
    import albumentations
except ImportError:
    albumentations = None

Number = Union[int, float]


@TRANSFORMS.register_module()
class GetBBoxCenterScale(BaseTransform):
    """Convert bboxes from [x, y, w, h] to center and scale.

    The center is the coordinates of the bbox center, and the scale is the
    bbox width and height normalized by a scale factor.

    Required Keys:

        - bbox

    Added Keys:

        - bbox_center
        - bbox_scale

    Args:
        padding (float): The bbox padding scale that will be multilied to
            `bbox_scale`. Defaults to 1.25
    """

    def __init__(self, padding: float = 1.25) -> None:
        super().__init__()

        self.padding = padding

    def transform(self, results: Dict) -> Optional[dict]:
        """The transform function of :class:`GetBBoxCenterScale`.

        See ``transform()`` method of :class:`BaseTransform` for details.

        Args:
            results (dict): The result dict

        Returns:
            dict: The result dict.
        """
        
        # Save the original bbox wrt. input
        results['bbox_xyxy_wrt_input'] = results['bbox']
        
        if 'bbox_center' in results and 'bbox_scale' in results:
            rank, _ = get_dist_info()
            if rank == 0:
                warnings.warn('Use the existing "bbox_center" and "bbox_scale"'
                              '. The padding will still be applied.')
            results['bbox_scale'] = results['bbox_scale'] * self.padding

        else:
            bbox = results['bbox']
            center, scale = bbox_xyxy2cs(bbox, padding=self.padding)

            results['bbox_center'] = center
            results['bbox_scale'] = scale

        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__ + f'(padding={self.padding})'
        return repr_str


@TRANSFORMS.register_module()
class MaskBackground(BaseTransform):
    """Convert bboxes from [x, y, w, h] to center and scale.

    The center is the coordinates of the bbox center, and the scale is the
    bbox width and height normalized by a scale factor.

    Required Keys:

        - bbox

    Added Keys:

        - bbox_center
        - bbox_scale

    Args:
        padding (float): The bbox padding scale that will be multilied to
            `bbox_scale`. Defaults to 1.25
    """

    def __init__(self,
        continue_on_failure: bool = True,
        prob: float = 1.0,
        alpha: float = 1.0,
        erode_prob: float = 0.0,
        erode_amount: float = 0.5,
        dilate_prob: float = 0.0,
        dilate_amount: float = 0.5,
    ) -> None:
        
        super().__init__()
        
        assert 0 <= alpha <= 1, 'alpha should be in [0, 1]'
        assert 0 <= prob <= 1, 'prob should be in [0, 1]'

        self.continue_on_failure = continue_on_failure
        self.alpha = alpha
        self.prob = prob
        
        assert 0 <= erode_prob <= 1, 'erode_prob should be in [0, 1]'
        assert 0 <= dilate_prob <= 1, 'dilate_prob should be in [0, 1]'
        assert 0 < erode_amount < 1, 'erode_amount should be in [0, 1]'
        assert 0 < dilate_amount < 1, 'dilate_amount should be in [0, 1]'
        assert erode_prob + dilate_prob <= 1, 'erode_prob + dilate_prob should be less than or equal to 1'
        self.noise_prob = erode_prob + dilate_prob
        if self.noise_prob > 0:
            self.erode_prob = erode_prob / (self.noise_prob)
            self.dilate_prob = dilate_prob / (self.noise_prob)
        else:
            self.erode_prob = 0
            self.dilate_prob = 0

        self.erode_amount = erode_amount
        self.dilate_amount = dilate_amount

    def _perturb_by_dilation(self, mask: np.ndarray) -> np.ndarray:
        """Perturb the mask to simulate real-world detector."""
        mask_shape = mask.shape

        mask_area = (mask>0).sum()

        # Close the mask to erase small holes
        k = max(mask_area // 1000, 5)
        kernel = np.ones((k, k), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

        # Dilate the mask to increase it a bit
        k = max(mask_area // 3000, 5)
        kernel = np.ones((k, k), np.uint8)
        mask = cv2.dilate(mask, kernel, iterations=1)
                    
        return mask.reshape(mask_shape)

    def _perturb_by_erosion(self, mask: np.ndarray) -> np.ndarray:
        """Perturb the mask to simulate real-world detector."""
        mask_shape = mask.shape

        mask_area = (mask>0).sum()

        # Close the mask to erase small holes
        k = max(mask_area // 1000, 5)
        kernel = np.ones((k, k), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

        # Erode the mask to decrease it a bit and cut-off limbs
        k = max(mask_area // 3000, 5)
        kernel = np.ones((k, k), np.uint8)
        mask = cv2.erode(mask, kernel, iterations=1)
                    
        return mask.reshape(mask_shape)

    @cache_randomness
    def _perturb_by_patches(self, mask: np.ndarray, amount: float, num_patches: int = 10) -> np.ndarray:
        mask_shape = mask.shape
        
        # Generate 10 random seeds uniformly distributed in the mask
        mask_idx = np.where(mask.flatten() > 0)[0]
        seeds = np.random.choice(mask_idx, num_patches, replace=False)
        sx, sy = np.unravel_index(seeds, mask.shape)

        # For each pixel, label it by it nearest seed
        labels = np.ones_like(mask)
        seed_labels = np.zeros_like(mask)
        seed_labels[sx, sy] = np.arange(num_patches) + 1

        _, indices = distance_transform_edt(seed_labels == 0, return_indices=True)
        labels = seed_labels[indices[0], indices[1]]
        labels = labels * mask

        # Select labels for removal
        random_remove_amount = np.random.uniform(0.0, amount)
        random_remove_ratio = int(num_patches * random_remove_amount)
        remove_labels = np.random.choice(np.unique(labels), random_remove_ratio, replace=False)
        binary_labels = np.isin(labels, remove_labels, invert=True)

        mask = (binary_labels > 0).astype(np.uint8) * mask

        return mask.reshape(mask_shape)

    @cache_randomness
    def _coin_flip(self) -> bool:
        return np.random.rand() < 0.5

    @cache_randomness
    def _perturb_mask(self, mask: np.ndarray) -> np.ndarray:
        """Perturb the mask to simulate real-world detector."""

        mask_shape = mask.shape

        if not np.random.rand() < self.noise_prob:
            return mask

        # Erode and dilate the mask to increase smoothness
        kernel = np.ones((5, 5), np.uint8)       
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        
        increase_mask = np.random.choice([False, True], p=[self.erode_prob, self.dilate_prob])
        
        if increase_mask:
            if self._coin_flip():
                try:
                    mask = self._perturb_by_patches(
                        mask=1-mask,
                        amount=self.dilate_amount,
                        num_patches=50,
                    )
                    mask = 1-mask
                except ValueError:
                    pass
            else:
                mask = self._perturb_by_dilation(mask)

        else:
            if self._coin_flip():
                try:
                    mask = self._perturb_by_patches(
                        mask=mask,
                        amount=self.erode_amount,
                        num_patches=10,
                    )
                except ValueError:
                    pass

            else:
                mask = self._perturb_by_erosion(mask)

        mask = (mask>0).astype(np.uint8)
        return mask.reshape(mask_shape)

    @cache_randomness
    def _do_masking(self):
        return np.random.rand() < self.prob

    def transform(self, results: Dict) -> Optional[dict]:
        """The transform function of :class:`GetBBoxCenterScale`.

        See ``transform()`` method of :class:`BaseTransform` for details.

        Args:
            results (dict): The result dict

        Returns:
            dict: The result dict.
        """

        # Try to load the mask from the results        
        mask = results.get('segmentation', None)
        # print("\nMaskBackground: ", mask is not None)

        if mask is None and not self.continue_on_failure:
            raise ValueError('No mask found in the results and self.continue_on_failure is set to False.')

        if mask is not None and self._do_masking():
            # Convert mask from polygons to binary mask
            try:            
                mask_rle = Mask.frPyObjects(mask, results['img_shape'][0], results['img_shape'][1])
            except IndexError:
                # breakpoint()
                # print("Mask shape:", mask.shape)
                # print("Mask max:", mask.max())
                # print("Mask min:", mask.min())
                # print("Image shape:", results['img_shape'])

                return results

            
            mask_rle = Mask.merge(mask_rle)
            img = results['img'].copy()
            masked_image = results['img'].copy()
            mask = Mask.decode(mask_rle).reshape((img.shape[0], img.shape[1], 1))
            binary_mask = (mask > 0).astype(np.uint8)

            # Perturb the mask to simulate real-world detector
            # print("Here I would perturb the mask")
            old_mask = mask.copy()
            binary_mask = self._perturb_mask(binary_mask)

            masked_image = masked_image * binary_mask
            results['img'] = cv2.addWeighted(img, 1 - self.alpha, masked_image, self.alpha, 0)

            # hash_id = abs(hash(555))
            # cv2.imwrite("tmp_visualization/_perturbed_mask_{:d}.jpg".format(hash_id), mask * 255)
            # cv2.imwrite("tmp_visualization/_old_mask_{:d}.jpg".format(hash_id), old_mask * 255)
            # cv2.imwrite("tmp_visualization/_weighted_masked_image_{:d}.jpg".format(hash_id), results['img'])
            # breakpoint()
            # Save the mask as a binary mask

        # Save the image
        img = results['img']
        
        # img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        # cv2.imwrite("tmp_visualization/masked_image_{:d}.jpg".format(abs(hash(555))), img)

        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__ + f'(continue_on_failure={self.continue_on_failure})'
        return repr_str


@TRANSFORMS.register_module()
class RandomFlip(BaseTransform):
    """Randomly flip the image, bbox and keypoints.

    Required Keys:

        - img
        - img_shape
        - flip_indices
        - input_size (optional)
        - bbox (optional)
        - bbox_center (optional)
        - keypoints (optional)
        - keypoints_visible (optional)
        - keypoints_visibility (optional)
        - img_mask (optional)

    Modified Keys:

        - img
        - bbox (optional)
        - bbox_center (optional)
        - keypoints (optional)
        - keypoints_visible (optional)
        - keypoints_visibility (optional)
        - img_mask (optional)

    Added Keys:

        - flip
        - flip_direction

    Args:
        prob (float | list[float]): The flipping probability. If a list is
            given, the argument `direction` should be a list with the same
            length. And each element in `prob` indicates the flipping
            probability of the corresponding one in ``direction``. Defaults
            to 0.5
        direction (str | list[str]): The flipping direction. Options are
            ``'horizontal'``, ``'vertical'`` and ``'diagonal'``. If a list is
            is given, each data sample's flipping direction will be sampled
            from a distribution determined by the argument ``prob``. Defaults
            to ``'horizontal'``.
    """

    def __init__(self,
                 prob: Union[float, List[float]] = 0.5,
                 direction: Union[str, List[str]] = 'horizontal') -> None:
        if isinstance(prob, list):
            assert is_list_of(prob, float)
            assert 0 <= sum(prob) <= 1
        elif isinstance(prob, float):
            assert 0 <= prob <= 1
        else:
            raise ValueError(f'probs must be float or list of float, but \
                              got `{type(prob)}`.')
        self.prob = prob

        valid_directions = ['horizontal', 'vertical', 'diagonal']
        if isinstance(direction, str):
            assert direction in valid_directions
        elif isinstance(direction, list):
            assert is_list_of(direction, str)
            assert set(direction).issubset(set(valid_directions))
        else:
            raise ValueError(f'direction must be either str or list of str, \
                               but got `{type(direction)}`.')
        self.direction = direction

        if isinstance(prob, list):
            assert len(prob) == len(self.direction)

    @cache_randomness
    def _choose_direction(self) -> str:
        """Choose the flip direction according to `prob` and `direction`"""
        if isinstance(self.direction,
                      List) and not isinstance(self.direction, str):
            # None means non-flip
            direction_list: list = list(self.direction) + [None]
        elif isinstance(self.direction, str):
            # None means non-flip
            direction_list = [self.direction, None]

        if isinstance(self.prob, list):
            non_prob: float = 1 - sum(self.prob)
            prob_list = self.prob + [non_prob]
        elif isinstance(self.prob, float):
            non_prob = 1. - self.prob
            # exclude non-flip
            single_ratio = self.prob / (len(direction_list) - 1)
            prob_list = [single_ratio] * (len(direction_list) - 1) + [non_prob]

        cur_dir = np.random.choice(direction_list, p=prob_list)

        return cur_dir

    def transform(self, results: dict) -> dict:
        """The transform function of :class:`RandomFlip`.

        See ``transform()`` method of :class:`BaseTransform` for details.

        Args:
            results (dict): The result dict

        Returns:
            dict: The result dict.
        """

        flip_dir = self._choose_direction()

        if flip_dir is None:
            results['flip'] = False
            results['flip_direction'] = None
        else:
            results['flip'] = True
            results['flip_direction'] = flip_dir

            h, w = results.get('input_size', results['img_shape'])
            # flip image and mask
            if isinstance(results['img'], list):
                results['img'] = [
                    imflip(img, direction=flip_dir) for img in results['img']
                ]
            else:
                results['img'] = imflip(results['img'], direction=flip_dir)

            if 'img_mask' in results:
                results['img_mask'] = imflip(
                    results['img_mask'], direction=flip_dir)

            # flip bboxes
            if results.get('bbox', None) is not None:
                results['bbox'] = flip_bbox(
                    results['bbox'],
                    image_size=(w, h),
                    bbox_format='xyxy',
                    direction=flip_dir)
            
            # flip bboxes
            if results.get('bbox_xyxy_wrt_input', None) is not None:
                results['bbox_xyxy_wrt_input'] = flip_bbox(
                    results['bbox_xyxy_wrt_input'],
                    image_size=(w, h),
                    bbox_format='xyxy',
                    direction=flip_dir)

            if results.get('bbox_center', None) is not None:
                results['bbox_center'] = flip_bbox(
                    results['bbox_center'],
                    image_size=(w, h),
                    bbox_format='center',
                    direction=flip_dir)

            # flip keypoints
            if results.get('keypoints', None) is not None:
                keypoints, keypoints_visible = flip_keypoints(
                    results['keypoints'],
                    results.get('keypoints_visible', None),
                    image_size=(w, h),
                    flip_indices=results['flip_indices'],
                    direction=flip_dir)
                _, keypoints_visibility = flip_keypoints(
                    results['keypoints'],
                    results.get('keypoints_visibility', None),
                    image_size=(w, h),
                    flip_indices=results['flip_indices'],
                    direction=flip_dir)

                results['keypoints'] = keypoints
                results['keypoints_visible'] = keypoints_visible
                results['keypoints_visibility'] = keypoints_visibility

        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__
        repr_str += f'(prob={self.prob}, '
        repr_str += f'direction={self.direction})'
        return repr_str


@TRANSFORMS.register_module()
class RandomHalfBody(BaseTransform):
    """Data augmentation with half-body transform that keeps only the upper or
    lower body at random.

    Required Keys:

        - keypoints
        - keypoints_visible
        - upper_body_ids
        - lower_body_ids

    Modified Keys:

        - bbox
        - bbox_center
        - bbox_scale

    Args:
        min_total_keypoints (int): The minimum required number of total valid
            keypoints of a person to apply half-body transform. Defaults to 8
        min_half_keypoints (int): The minimum required number of valid
            half-body keypoints of a person to apply half-body transform.
            Defaults to 2
        padding (float): The bbox padding scale that will be multilied to
            `bbox_scale`. Defaults to 1.5
        prob (float): The probability to apply half-body transform when the
            keypoint number meets the requirement. Defaults to 0.3
    """

    def __init__(self,
                 min_total_keypoints: int = 9,
                 min_upper_keypoints: int = 2,
                 min_lower_keypoints: int = 3,
                 padding: float = 1.5,
                 prob: float = 0.3,
                 upper_prioritized_prob: float = 0.7) -> None:
        super().__init__()
        self.min_total_keypoints = min_total_keypoints
        self.min_upper_keypoints = min_upper_keypoints
        self.min_lower_keypoints = min_lower_keypoints
        self.padding = padding
        self.prob = prob
        self.upper_prioritized_prob = upper_prioritized_prob

    def _get_half_body_bbox(self, keypoints: np.ndarray,
                            half_body_ids: List[int]
                            ) -> Tuple[np.ndarray, np.ndarray]:
        """Get half-body bbox center and scale of a single instance.

        Args:
            keypoints (np.ndarray): Keypoints in shape (K, D)
            upper_body_ids (list): The list of half-body keypont indices

        Returns:
            tuple: A tuple containing half-body bbox center and scale
            - center: Center (x, y) of the bbox
            - scale: Scale (w, h) of the bbox
        """

        selected_keypoints = keypoints[half_body_ids]
        center = selected_keypoints.mean(axis=0)[:2]

        x1, y1 = selected_keypoints.min(axis=0)
        x2, y2 = selected_keypoints.max(axis=0)
        w = x2 - x1
        h = y2 - y1
        scale = np.array([w, h], dtype=center.dtype) * self.padding

        return center, scale
    
    def _get_half_body_exact_bbox(self, keypoints: np.ndarray,
                                half_body_ids: List[int],
                                bbox: np.ndarray,
                                ) -> np.ndarray:
        """Get half-body bbox center and scale of a single instance.

        Args:
            keypoints (np.ndarray): Keypoints in shape (K, D)
            upper_body_ids (list): The list of half-body keypont indices

        Returns:
            tuple: A tuple containing half-body bbox center and scale
            - center: Center (x, y) of the bbox
            - scale: Scale (w, h) of the bbox
        """

        selected_keypoints = keypoints[half_body_ids]
        center = selected_keypoints.mean(axis=0)[:2]

        x1, y1 = selected_keypoints.min(axis=0)
        x2, y2 = selected_keypoints.max(axis=0)
        w = x2 - x1
        h = y2 - y1
        scale = np.array([w, h], dtype=center.dtype) * self.padding

        x1, y1 = center - scale / 2
        x2, y2 = center + scale / 2

        # Do not exceed the original bbox
        x1 = np.maximum(x1, bbox[0])
        y1 = np.maximum(y1, bbox[1])
        x2 = np.minimum(x2, bbox[2])
        y2 = np.minimum(y2, bbox[3])

        return np.array([x1, y1, x2, y2])

    @cache_randomness
    def _random_select_half_body(self, keypoints_visible: np.ndarray,
                                 upper_body_ids: List[int],
                                 lower_body_ids: List[int]
                                 ) -> List[Optional[List[int]]]:
        """Randomly determine whether applying half-body transform and get the
        half-body keyponit indices of each instances.

        Args:
            keypoints_visible (np.ndarray, optional): The visibility of
                keypoints in shape (N, K, 1) or (N, K, 2).
            upper_body_ids (list): The list of upper body keypoint indices
            lower_body_ids (list): The list of lower body keypoint indices

        Returns:
            list[list[int] | None]: The selected half-body keypoint indices
            of each instance. ``None`` means not applying half-body transform.
        """

        if keypoints_visible.ndim == 3:
            keypoints_visible = keypoints_visible[..., 0]

        half_body_ids = []

        for visible in keypoints_visible:
            if visible.sum() < self.min_total_keypoints:
                indices = None
            elif np.random.rand() > self.prob:
                indices = None
            else:
                upper_valid_ids = [i for i in upper_body_ids if visible[i] > 0]
                lower_valid_ids = [i for i in lower_body_ids if visible[i] > 0]

                num_upper = len(upper_valid_ids)
                num_lower = len(lower_valid_ids)

                prefer_upper = np.random.rand() < self.upper_prioritized_prob
                if (num_upper < self.min_upper_keypoints
                        and num_lower < self.min_lower_keypoints):
                    indices = None
                elif num_lower < self.min_lower_keypoints:
                    indices = upper_valid_ids
                elif num_upper < self.min_upper_keypoints:
                    indices = lower_valid_ids
                else:
                    indices = (
                        upper_valid_ids if prefer_upper else lower_valid_ids)

            half_body_ids.append(indices)

        return half_body_ids

    def transform(self, results: Dict) -> Optional[dict]:
        """The transform function of :class:`HalfBodyTransform`.

        See ``transform()`` method of :class:`BaseTransform` for details.

        Args:
            results (dict): The result dict

        Returns:
            dict: The result dict.
        """
        half_body_ids = self._random_select_half_body(
            keypoints_visible=results['keypoints_visible'],
            upper_body_ids=results['upper_body_ids'],
            lower_body_ids=results['lower_body_ids'])

        bbox_center = []
        bbox_scale = []

        bbox_xyxy_wrt_input = []

        for i, indices in enumerate(half_body_ids):
            if indices is None:
                bbox_center.append(results['bbox_center'][i])
                bbox_scale.append(results['bbox_scale'][i])
                bbox_xyxy_wrt_input.append(results['bbox_xyxy_wrt_input'][i])
            else:
                _center, _scale = self._get_half_body_bbox(
                    results['keypoints'][i], indices)
                bbox_center.append(_center)
                bbox_scale.append(_scale)
                exact_bbox = self._get_half_body_exact_bbox(
                    results['keypoints'][i], indices, results['bbox_xyxy_wrt_input'][i])
                bbox_xyxy_wrt_input.append(exact_bbox)

        results['bbox_center'] = np.stack(bbox_center)
        results['bbox_scale'] = np.stack(bbox_scale)
        results['bbox_xyxy_wrt_input'] = np.stack(bbox_xyxy_wrt_input)
        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__
        repr_str += f'(min_total_keypoints={self.min_total_keypoints}, '
        repr_str += f'min_upper_keypoints={self.min_upper_keypoints}, '
        repr_str += f'min_lower_keypoints={self.min_lower_keypoints}, '
        repr_str += f'padding={self.padding}, '
        repr_str += f'prob={self.prob}, '
        repr_str += f'upper_prioritized_prob={self.upper_prioritized_prob})'
        return repr_str


@TRANSFORMS.register_module()
class RandomPatchesBlackout(BaseTransform):
    """Data augmentation that divide image into patches and set color of random
        pathes to black. In AID paper marked as 'hide and seek'.

    Required Keys:

        - keypoints
        - keypoints_visible
        - keypoint_visibility

    Modified Keys:

        - img
        - keypoint_visibility

    Args:
        grid_size (tuple(int, int)): Grid size of the patches. Defaults to
            (8, 6)
        mask_ratio (float): Ratio of patches to blackout. Defaults to 0.3
        prob (float): The probability to apply black patches. Defaults to 0.8
    """

    def __init__(self,
                 grid_size: Tuple[int, int] = (8, 6),
                 mask_ratio: float = 0.3,
                 prob: float = 0.8) -> None:
        super().__init__()
        self.grid_size = grid_size
        self.mask_ratio = mask_ratio
        self.prob = prob

    @cache_randomness
    def _get_random_patches(self, grid_h, grid_w) -> np.ndarray:
        black_patches = np.zeros((grid_h, grid_w), dtype=bool)

        if np.random.rand() < self.prob:
        
            # Split image into grid
            num_patches = int(self.grid_size[0] * self.grid_size[1])

            # Randomly choose patches to blackout
            black_patches = np.random.choice(
                [0, 1],
                num_patches,
                p=[1 - self.mask_ratio, self.mask_ratio]
            )
            black_patches = black_patches.reshape(grid_h, grid_w).astype(bool)

        return black_patches


    def transform(self, results: Dict) -> Optional[dict]:
        """The transform function of :class:`HalfBodyTransform`.

        See ``transform()`` method of :class:`BaseTransform` for details.

        Args:
            results (dict): The result dict

        Returns:
            dict: The result dict.
        """

        img = results['img']
        
        if "transformed_keypoints" in results:
            kpts = results['transformed_keypoints'].squeeze()
        else:
            kpts = results['keypoints'].squeeze()
        
        h, w = img.shape[:2]
        grid_h, grid_w = self.grid_size
        dh = np.ceil(h / grid_h).astype(int)
        dw = np.ceil(w / grid_w).astype(int)

        black_patches = self._get_random_patches(grid_h, grid_w)

        for i in range(grid_h):
            for j in range(grid_w):
                if black_patches[i, j]:
                    # Set all pixel in the patch to black
                    img[i*dh : (i+1)*dh, j*dw : (j+1)*dw, :] = 0


                    # Set keypoints in the patch to invisible
                    in_black = (
                        (kpts[:, 0] >= j*dw) &
                        (kpts[:, 0] < (j+1)*dw) &
                        (kpts[:, 1] >= i*dh) &
                        (kpts[:, 1] < (i+1)*dh)
                    )
                    results['keypoints_visibility'][:, in_black] = 0
                        
        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__
        repr_str += f'(grid_size={self.grid_size}, '
        repr_str += f'mask_ratio={self.mask_ratio}, '
        repr_str += f'prob={self.prob})'
        return repr_str


@TRANSFORMS.register_module()
class RandomEdgesBlackout(BaseTransform):
    """Data augmentation that masks edged of the image with black color
        simulating image edge or random texture.

    Required Keys:

        - keypoints
        - keypoints_visible
        - keypoint_visibility

    Modified Keys:

        - img
        - keypoint_visibility

    Args:
        mask_ratio_range (tuple[float, float]): Range or mask-to-image ratio. Defaults to
            (0.1, 0.3)
        prob (float): The probability to apply black patches. Defaults to 0.8
        texture_prob (float): The probability to apply texture to the blackout area. Defaults to 0.0
    """

    def __init__(self,
                 mask_ratio_range: tuple[float, float] = (0.1, 0.3),
                 prob: float = 0.8,
                 texture_prob: float = 0.0,
                 context_size:float = 1.25) -> None:
        super().__init__()
        self.mask_ratio_range = mask_ratio_range
        self.prob = prob
        self.texture_prob = texture_prob
        self.context_size = context_size

    @cache_randomness
    def _get_random_mask(self, w, h, bbox_xyxy) -> float:
        """Get random mask ratio.

        Args:
            w (int): Width of the image
            h (int): Height of the image
            bbox_xyxy (tuple): Bounding box (x1, y1, x2, y2)

        Returns:
            np.array: mask (1 for blackout, 0 for keep)
            tuple: bounds of the blackout area (x1, y1, x2, y2)
        """
        mask = np.zeros((h, w), dtype=bool)
        bbox_c, bbox_s = bbox_xyxy2cs(bbox_xyxy, padding=self.context_size)
        x0, y0, x1, y1 = bbox_cs2xyxy(bbox_c, bbox_s)

        # Clip the bounding box to the image
        x0 = np.maximum(x0, 0).astype(int)
        y0 = np.maximum(y0, 0).astype(int)
        x1 = np.minimum(x1, w).astype(int)
        y1 = np.minimum(y1, h).astype(int)
        
        # Set default values
        x = 0
        y = 0
        dw = w
        dh = h
        is_textured = False

        if np.random.rand() < self.prob:
            
            # Generate random rectangle to keep
            rh, rw = np.random.uniform(
                1-self.mask_ratio_range[1],
                1-self.mask_ratio_range[0],
                2
            )
            dh = int((y1-y0) * rh)
            dw = int((x1-x0) * rw)
            x_end = x1-dw if x1-dw > x0 else x0+1
            y_end = y1-dh if y1-dh > y0 else y0+1
            try:
                x = np.random.randint(x0, x_end)
                y = np.random.randint(y0, y_end)
            except ValueError:
                print(x, x0, dw, x1, x1-dw, x_end)
                print(y, y0, dh, y1, y1-dh, y_end)
                raise ValueError

            # Set all pixel outside of the rectangle to black
            mask[y:y+dh, x:x+dw] = True
            
            # Invert the mask. True means blackout
            mask = ~mask

            # Add texture
            is_textured = np.random.rand() < self.texture_prob

        return mask, (x, y, dw+x, dh+y), is_textured

    def _get_random_color(self) -> np.ndarray:
        """Get random color.

        Returns:
            np.array: color
        """
        h = np.random.randint(0, 360)
        s = np.random.uniform(0.75, 1)
        l = np.random.uniform(0.3, 0.7)
        hls_color = np.array([h, l, s])
        rgb_color = cv2.cvtColor(
            np.array([[hls_color]], dtype=np.float32),
            cv2.COLOR_HLS2RGB
        ).squeeze() * 255
        color = rgb_color.astype(np.uint8)
        return color.tolist()

    def _get_random_texture(self, w, h) -> np.ndarray:
        """Get random texture.

        Args:
            w (int): Width of the image
            h (int): Height of the image

        Returns:
            np.array: texture
        """
        mode = np.random.choice([
            'lines',
            'squares',
            'circles',
            # 'noise',
            # 'uniform',
        ])

        if mode == 'lines':
            texture = np.zeros((h, w, 3), dtype=np.uint8)
            texture[:, :, :] = self._get_random_color()
            num_lines = np.random.randint(1, 20)
            for _ in range(num_lines):
                x1, y1 = np.random.randint(0, w), np.random.randint(0, h)
                x2, y2 = np.random.randint(0, w), np.random.randint(0, h)
                line_width = np.random.randint(1, 10)
                color = self._get_random_color()
                cv2.line(texture, (x1, y1), (x2, y2), color, line_width)
        elif mode == 'squares':
            texture = np.zeros((h, w, 3), dtype=np.uint8)
            texture[:, :, :] = self._get_random_color()
            num_squares = np.random.randint(1, 20)
            for _ in range(num_squares):
                x1, y1 = np.random.randint(0, w), np.random.randint(0, h)
                x2, y2 = np.random.randint(0, w), np.random.randint(0, h)
                color = self._get_random_color()
                cv2.rectangle(texture, (x1, y1), (x2, y2), color, -1)
        elif mode == 'circles':
            texture = np.zeros((h, w, 3), dtype=np.uint8)
            texture[:, :, :] = self._get_random_color()
            num_circles = np.random.randint(1, 20)
            for _ in range(num_circles):
                x, y = np.random.randint(0, w), np.random.randint(0, h)
                r = np.random.randint(1, min(w, h) // 2)
                color = self._get_random_color()
                cv2.circle(texture, (x, y), r, color, -1)
        elif mode == 'noise':
            texture = np.random.randint(0, 256, (h, w, 3), dtype=np.uint8)
        elif mode == 'uniform':
            texture = np.zeros((h, w, 3), dtype=np.uint8)
            texture[:, :, :] = self._get_random_color()

        return texture

    def transform(self, results: Dict) -> Optional[dict]:
        """The transform function of :class:`HalfBodyTransform`.

        See ``transform()`` method of :class:`BaseTransform` for details.

        Args:
            results (dict): The result dict

        Returns:
            dict: The result dict.
        """

        img = results['img']
        
        if "transformed_keypoints" in results:
            kpts = results['transformed_keypoints'].squeeze()
        else:
            kpts = results['keypoints'].squeeze()

        # Generate random mask
        mask, (x1, y1, x2, y2), is_textured = self._get_random_mask(img.shape[1], img.shape[0], results['bbox_xyxy_wrt_input'].flatten())
        # breakpoint()
        # print("img shape", img.shape)
        # print("results", results.keys())

        # Apply the mask
        if is_textured:
            textured_img = self._get_random_texture(img.shape[1], img.shape[0])
            textured_img[~mask, :] = img[~mask, :]
            img = textured_img
        else:
            # Set all pixel outside of the rectangle to black
            img[mask, :] = 0
        results['img'] = img

        # Set keypoints outside of the rectangle to invisible
        in_rect = (
            (kpts[:, 0] >= x1) &
            (kpts[:, 0] < x2) &
            (kpts[:, 1] >= y1) &
            (kpts[:, 1] < y2)
        )
        results['keypoints_visibility'][:, ~in_rect] = 0

        # Create new entry describing keypoints in the 'cropped' area
        results['keypoints_in_image'] = in_rect.squeeze().astype(int)

        # Crop the bbox_xyxy_wrt_input according to the blackout area
        if 'bbox_xyxy_wrt_input' in results:
            bbox_xyxy = results['bbox_xyxy_wrt_input'].flatten()
            bbox_xyxy[0] = np.maximum(bbox_xyxy[0], x1)
            bbox_xyxy[1] = np.maximum(bbox_xyxy[1], y1)
            bbox_xyxy[2] = np.minimum(bbox_xyxy[2], x2)
            bbox_xyxy[3] = np.minimum(bbox_xyxy[3], y2)
            results['bbox_xyxy_wrt_input'] = bbox_xyxy.reshape(-1, 4)

        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__
        repr_str += f'(mask_ratio_range={self.mask_ratio_range}, '
        repr_str += f'prob={self.prob}), '
        repr_str += f'texture_prob={self.texture_prob})'
        return repr_str


@TRANSFORMS.register_module()
class RandomBBoxTransform(BaseTransform):
    r"""Rnadomly shift, resize and rotate the bounding boxes.

    Required Keys:

        - bbox_center
        - bbox_scale

    Modified Keys:

        - bbox_center
        - bbox_scale

    Added Keys:
        - bbox_rotation

    Args:
        shift_factor (float): Randomly shift the bbox in range
            :math:`[-dx, dx]` and :math:`[-dy, dy]` in X and Y directions,
            where :math:`dx(y) = x(y)_scale \cdot shift_factor` in pixels.
            Defaults to 0.16
        shift_prob (float): Probability of applying random shift. Defaults to
            0.3
        scale_factor (Tuple[float, float]): Randomly resize the bbox in range
            :math:`[scale_factor[0], scale_factor[1]]`. Defaults to (0.5, 1.5)
        scale_prob (float): Probability of applying random resizing. Defaults
            to 1.0
        rotate_factor (float): Randomly rotate the bbox in
            :math:`[-rotate_factor, rotate_factor]` in degrees. Defaults
            to 80.0
        rotate_prob (float): Probability of applying random rotation. Defaults
            to 0.6
    """

    def __init__(self,
                 shift_factor: float = 0.16,
                 shift_prob: float = 0.3,
                 scale_factor: Tuple[float, float] = (0.5, 1.5),
                 scale_prob: float = 1.0,
                 rotate_factor: float = 80.0,
                 rotate_prob: float = 0.6) -> None:
        super().__init__()

        self.shift_factor = shift_factor
        self.shift_prob = shift_prob
        self.scale_factor = scale_factor
        self.scale_prob = scale_prob
        self.rotate_factor = rotate_factor
        self.rotate_prob = rotate_prob

    @staticmethod
    def _truncnorm(low: float = -1.,
                   high: float = 1.,
                   size: tuple = ()) -> np.ndarray:
        """Sample from a truncated normal distribution."""
        return truncnorm.rvs(low, high, size=size).astype(np.float32)

    @cache_randomness
    def _get_transform_params(self, num_bboxes: int) -> Tuple:
        """Get random transform parameters.

        Args:
            num_bboxes (int): The number of bboxes

        Returns:
            tuple:
            - offset (np.ndarray): Offset factor of each bbox in shape (n, 2)
            - scale (np.ndarray): Scaling factor of each bbox in shape (n, 1)
            - rotate (np.ndarray): Rotation degree of each bbox in shape (n,)
        """
        random_v = self._truncnorm(size=(num_bboxes, 4))
        offset_v = random_v[:, :2]
        scale_v = random_v[:, 2:3]
        rotate_v = random_v[:, 3]

        # Get shift parameters
        offset = offset_v * self.shift_factor
        offset = np.where(
            np.random.rand(num_bboxes, 1) < self.shift_prob, offset, 0.)

        # Get scaling parameters
        scale_min, scale_max = self.scale_factor
        mu = (scale_max + scale_min) * 0.5
        sigma = (scale_max - scale_min) * 0.5
        scale = scale_v * sigma + mu
        scale = np.where(
            np.random.rand(num_bboxes, 1) < self.scale_prob, scale, 1.)

        # Get rotation parameters
        rotate = rotate_v * self.rotate_factor
        rotate = np.where(
            np.random.rand(num_bboxes) < self.rotate_prob, rotate, 0.)

        return offset, scale, rotate

    def transform(self, results: Dict) -> Optional[dict]:
        """The transform function of :class:`RandomBboxTransform`.

        See ``transform()`` method of :class:`BaseTransform` for details.

        Args:
            results (dict): The result dict

        Returns:
            dict: The result dict.
        """
        bbox_scale = results['bbox_scale']
        num_bboxes = bbox_scale.shape[0]

        offset, scale, rotate = self._get_transform_params(num_bboxes)

        results['bbox_center'] = results['bbox_center'] + offset * bbox_scale
        results['bbox_scale'] = results['bbox_scale'] * scale
        results['bbox_rotation'] = rotate

        bbox_xyxy_wrt_input = results.get('bbox_xyxy_wrt_input', None)
        if bbox_xyxy_wrt_input is not None:
            _c, _s = bbox_xyxy2cs(bbox_xyxy_wrt_input, padding=1.0)
            _c = _c + offset * _s
            _s = _s * scale
            results['bbox_xyxy_wrt_input'] = bbox_cs2xyxy(_c, _s).flatten()

        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__
        repr_str += f'(shift_prob={self.shift_prob}, '
        repr_str += f'shift_factor={self.shift_factor}, '
        repr_str += f'scale_prob={self.scale_prob}, '
        repr_str += f'scale_factor={self.scale_factor}, '
        repr_str += f'rotate_prob={self.rotate_prob}, '
        repr_str += f'rotate_factor={self.rotate_factor})'
        return repr_str


@TRANSFORMS.register_module()
@avoid_cache_randomness
class Albumentation(BaseTransform):
    """Albumentation augmentation (pixel-level transforms only).

    Adds custom pixel-level transformations from Albumentations library.
    Please visit `https://albumentations.ai/docs/`
    to get more information.

    Note: we only support pixel-level transforms.
    Please visit `https://github.com/albumentations-team/`
    `albumentations#pixel-level-transforms`
    to get more information about pixel-level transforms.

    Required Keys:

    - img

    Modified Keys:

    - img

    Args:
        transforms (List[dict]): A list of Albumentation transforms.
            An example of ``transforms`` is as followed:
            .. code-block:: python

                [
                    dict(
                        type='RandomBrightnessContrast',
                        brightness_limit=[0.1, 0.3],
                        contrast_limit=[0.1, 0.3],
                        p=0.2),
                    dict(type='ChannelShuffle', p=0.1),
                    dict(
                        type='OneOf',
                        transforms=[
                            dict(type='Blur', blur_limit=3, p=1.0),
                            dict(type='MedianBlur', blur_limit=3, p=1.0)
                        ],
                        p=0.1),
                ]
        keymap (dict | None): key mapping from ``input key`` to
            ``albumentation-style key``.
            Defaults to None, which will use {'img': 'image'}.
    """

    def __init__(self,
                 transforms: List[dict],
                 keymap: Optional[dict] = None) -> None:
        if albumentations is None:
            raise RuntimeError('albumentations is not installed')

        self.transforms = transforms

        self.aug = albumentations.Compose(
            [self.albu_builder(t) for t in self.transforms])

        if not keymap:
            self.keymap_to_albu = {
                'img': 'image',
            }
        else:
            self.keymap_to_albu = keymap

    def albu_builder(self, cfg: dict) -> albumentations:
        """Import a module from albumentations.

        It resembles some of :func:`build_from_cfg` logic.

        Args:
            cfg (dict): Config dict. It should at least contain the key "type".

        Returns:
            albumentations.BasicTransform: The constructed transform object
        """

        assert isinstance(cfg, dict) and 'type' in cfg
        args = cfg.copy()

        obj_type = args.pop('type')
        if mmengine.is_str(obj_type):
            if albumentations is None:
                raise RuntimeError('albumentations is not installed')
            rank, _ = get_dist_info()
            if rank == 0 and not hasattr(
                    albumentations.augmentations.transforms, obj_type):
                warnings.warn(
                    f'{obj_type} is not pixel-level transformations. '
                    'Please use with caution.')
            obj_cls = getattr(albumentations, obj_type)
        elif isinstance(obj_type, type):
            obj_cls = obj_type
        else:
            raise TypeError(f'type must be a str, but got {type(obj_type)}')

        if 'transforms' in args:
            args['transforms'] = [
                self.albu_builder(transform)
                for transform in args['transforms']
            ]

        return obj_cls(**args)

    def transform(self, results: dict) -> dict:
        """The transform function of :class:`Albumentation` to apply
        albumentations transforms.

        See ``transform()`` method of :class:`BaseTransform` for details.

        Args:
            results (dict): Result dict from the data pipeline.

        Return:
            dict: updated result dict.
        """
        # map result dict to albumentations format
        results_albu = {}
        for k, v in self.keymap_to_albu.items():
            assert k in results, \
                f'The `{k}` is required to perform albumentations transforms'
            results_albu[v] = results[k]

        # Apply albumentations transforms
        results_albu = self.aug(**results_albu)

        # map the albu results back to the original format
        for k, v in self.keymap_to_albu.items():
            results[k] = results_albu[v]

        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__ + f'(transforms={self.transforms})'
        return repr_str


@TRANSFORMS.register_module()
class PhotometricDistortion(BaseTransform):
    """Apply photometric distortion to image sequentially, every transformation
    is applied with a probability of 0.5. The position of random contrast is in
    second or second to last.

    1. random brightness
    2. random contrast (mode 0)
    3. convert color from BGR to HSV
    4. random saturation
    5. random hue
    6. convert color from HSV to BGR
    7. random contrast (mode 1)
    8. randomly swap channels

    Required Keys:

    - img

    Modified Keys:

    - img

    Args:
        brightness_delta (int): delta of brightness.
        contrast_range (tuple): range of contrast.
        saturation_range (tuple): range of saturation.
        hue_delta (int): delta of hue.
    """

    def __init__(self,
                 brightness_delta: int = 32,
                 contrast_range: Sequence[Number] = (0.5, 1.5),
                 saturation_range: Sequence[Number] = (0.5, 1.5),
                 hue_delta: int = 18) -> None:
        self.brightness_delta = brightness_delta
        self.contrast_lower, self.contrast_upper = contrast_range
        self.saturation_lower, self.saturation_upper = saturation_range
        self.hue_delta = hue_delta

    @cache_randomness
    def _random_flags(self) -> Sequence[Number]:
        """Generate the random flags for subsequent transforms.

        Returns:
            Sequence[Number]: a sequence of numbers that indicate whether to
                do the corresponding transforms.
        """
        # contrast_mode == 0 --> do random contrast first
        # contrast_mode == 1 --> do random contrast last
        contrast_mode = np.random.randint(2)
        # whether to apply brightness distortion
        brightness_flag = np.random.randint(2)
        # whether to apply contrast distortion
        contrast_flag = np.random.randint(2)
        # the mode to convert color from BGR to HSV
        hsv_mode = np.random.randint(4)
        # whether to apply channel swap
        swap_flag = np.random.randint(2)

        # the beta in `self._convert` to be added to image array
        # in brightness distortion
        brightness_beta = np.random.uniform(-self.brightness_delta,
                                            self.brightness_delta)
        # the alpha in `self._convert` to be multiplied to image array
        # in contrast distortion
        contrast_alpha = np.random.uniform(self.contrast_lower,
                                           self.contrast_upper)
        # the alpha in `self._convert` to be multiplied to image array
        # in saturation distortion to hsv-formatted img
        saturation_alpha = np.random.uniform(self.saturation_lower,
                                             self.saturation_upper)
        # delta of hue to add to image array in hue distortion
        hue_delta = np.random.randint(-self.hue_delta, self.hue_delta)
        # the random permutation of channel order
        swap_channel_order = np.random.permutation(3)

        return (contrast_mode, brightness_flag, contrast_flag, hsv_mode,
                swap_flag, brightness_beta, contrast_alpha, saturation_alpha,
                hue_delta, swap_channel_order)

    def _convert(self,
                 img: np.ndarray,
                 alpha: float = 1,
                 beta: float = 0) -> np.ndarray:
        """Multiple with alpha and add beta with clip.

        Args:
            img (np.ndarray): The image array.
            alpha (float): The random multiplier.
            beta (float): The random offset.

        Returns:
            np.ndarray: The updated image array.
        """
        img = img.astype(np.float32) * alpha + beta
        img = np.clip(img, 0, 255)
        return img.astype(np.uint8)

    def transform(self, results: dict) -> dict:
        """The transform function of :class:`PhotometricDistortion` to perform
        photometric distortion on images.

        See ``transform()`` method of :class:`BaseTransform` for details.


        Args:
            results (dict): Result dict from the data pipeline.

        Returns:
            dict: Result dict with images distorted.
        """

        assert 'img' in results, '`img` is not found in results'
        img = results['img']

        (contrast_mode, brightness_flag, contrast_flag, hsv_mode, swap_flag,
         brightness_beta, contrast_alpha, saturation_alpha, hue_delta,
         swap_channel_order) = self._random_flags()

        # random brightness distortion
        if brightness_flag:
            img = self._convert(img, beta=brightness_beta)

        # contrast_mode == 0 --> do random contrast first
        # contrast_mode == 1 --> do random contrast last
        if contrast_mode == 1:
            if contrast_flag:
                img = self._convert(img, alpha=contrast_alpha)

        if hsv_mode:
            # random saturation/hue distortion
            img = mmcv.bgr2hsv(img)
            if hsv_mode == 1 or hsv_mode == 3:
                # apply saturation distortion to hsv-formatted img
                img[:, :, 1] = self._convert(
                    img[:, :, 1], alpha=saturation_alpha)
            if hsv_mode == 2 or hsv_mode == 3:
                # apply hue distortion to hsv-formatted img
                img[:, :, 0] = img[:, :, 0].astype(int) + hue_delta
            img = mmcv.hsv2bgr(img)

        if contrast_mode == 1:
            if contrast_flag:
                img = self._convert(img, alpha=contrast_alpha)

        # randomly swap channels
        if swap_flag:
            img = img[..., swap_channel_order]

        results['img'] = img
        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__
        repr_str += (f'(brightness_delta={self.brightness_delta}, '
                     f'contrast_range=({self.contrast_lower}, '
                     f'{self.contrast_upper}), '
                     f'saturation_range=({self.saturation_lower}, '
                     f'{self.saturation_upper}), '
                     f'hue_delta={self.hue_delta})')
        return repr_str


@TRANSFORMS.register_module()
class GenerateTarget(BaseTransform):
    """Encode keypoints into Target.

    The generated target is usually the supervision signal of the model
    learning, e.g. heatmaps or regression labels.

    Required Keys:

        - keypoints
        - keypoints_visible
        - dataset_keypoint_weights

    Added Keys:

        - The keys of the encoded items from the codec will be updated into
            the results, e.g. ``'heatmaps'`` or ``'keypoint_weights'``. See
            the specific codec for more details.

    Args:
        encoder (dict | list[dict]): The codec config for keypoint encoding.
            Both single encoder and multiple encoders (given as a list) are
            supported
        multilevel (bool): Determine the method to handle multiple encoders.
            If ``multilevel==True``, generate multilevel targets from a group
            of encoders of the same type (e.g. multiple :class:`MSRAHeatmap`
            encoders with different sigma values); If ``multilevel==False``,
            generate combined targets from a group of different encoders. This
            argument will have no effect in case of single encoder. Defaults
            to ``False``
        use_dataset_keypoint_weights (bool): Whether use the keypoint weights
            from the dataset meta information. Defaults to ``False``
        target_type (str, deprecated): This argument is deprecated and has no
            effect. Defaults to ``None``
    """

    def __init__(self,
                 encoder: MultiConfig,
                 target_type: Optional[str] = None,
                 multilevel: bool = False,
                 use_dataset_keypoint_weights: bool = False) -> None:
        super().__init__()

        if target_type is not None:
            rank, _ = get_dist_info()
            if rank == 0:
                warnings.warn(
                    'The argument `target_type` is deprecated in'
                    ' GenerateTarget. The target type and encoded '
                    'keys will be determined by encoder(s).',
                    DeprecationWarning)

        self.encoder_cfg = deepcopy(encoder)
        self.multilevel = multilevel
        self.use_dataset_keypoint_weights = use_dataset_keypoint_weights

        if isinstance(self.encoder_cfg, list):
            self.encoder = [
                KEYPOINT_CODECS.build(cfg) for cfg in self.encoder_cfg
            ]
        else:
            assert not self.multilevel, (
                'Need multiple encoder configs if ``multilevel==True``')
            self.encoder = KEYPOINT_CODECS.build(self.encoder_cfg)

    def transform(self, results: Dict) -> Optional[dict]:
        """The transform function of :class:`GenerateTarget`.

        See ``transform()`` method of :class:`BaseTransform` for details.
        """

        if results.get('transformed_keypoints', None) is not None:
            # use keypoints transformed by TopdownAffine
            keypoints = results['transformed_keypoints']
        elif results.get('keypoints', None) is not None:
            # use original keypoints
            keypoints = results['keypoints']
        else:
            raise ValueError(
                'GenerateTarget requires \'transformed_keypoints\' or'
                ' \'keypoints\' in the results.')

        keypoints_visible = results['keypoints_visible']
        if keypoints_visible.ndim == 3 and keypoints_visible.shape[2] == 2:
            keypoints_visible, keypoints_visible_weights = \
                keypoints_visible[..., 0], keypoints_visible[..., 1]
            results['keypoints_visible'] = keypoints_visible
            results['keypoints_visible_weights'] = keypoints_visible_weights
        
        id_similarity = results.get('id_similarity', np.array([0]))
        keypoints_visibility = results.get("keypoints_visibility", None)    
    
        # Encoded items from the encoder(s) will be updated into the results.
        # Please refer to the document of the specific codec for details about
        # encoded items.
        if not isinstance(self.encoder, list):
            # For single encoding, the encoded items will be directly added
            # into results.
            auxiliary_encode_kwargs = {
                key: results[key]
                for key in self.encoder.auxiliary_encode_keys
            }
            encoded = self.encoder.encode(
                keypoints=keypoints,
                keypoints_visible=keypoints_visible,
                keypoints_visibility=keypoints_visibility,
                id_similarity=id_similarity,
                **auxiliary_encode_kwargs)

            if self.encoder.field_mapping_table:
                encoded[
                    'field_mapping_table'] = self.encoder.field_mapping_table
            if self.encoder.instance_mapping_table:
                encoded['instance_mapping_table'] = \
                    self.encoder.instance_mapping_table
            if self.encoder.label_mapping_table:
                encoded[
                    'label_mapping_table'] = self.encoder.label_mapping_table

        else:
            encoded_list = []
            _field_mapping_table = dict()
            _instance_mapping_table = dict()
            _label_mapping_table = dict()
            for _encoder in self.encoder:
                auxiliary_encode_kwargs = {
                    key: results[key]
                    for key in _encoder.auxiliary_encode_keys
                }
                encoded_list.append(
                    _encoder.encode(
                        keypoints=keypoints,
                        keypoints_visible=keypoints_visible,
                        keypoints_visibility=keypoints_visibility,
                        id_similarity=id_similarity,
                        **auxiliary_encode_kwargs))

                _field_mapping_table.update(_encoder.field_mapping_table)
                _instance_mapping_table.update(_encoder.instance_mapping_table)
                _label_mapping_table.update(_encoder.label_mapping_table)

            if self.multilevel:
                # For multilevel encoding, the encoded items from each encoder
                # should have the same keys.

                keys = encoded_list[0].keys()
                if not all(_encoded.keys() == keys
                           for _encoded in encoded_list):
                    raise ValueError(
                        'Encoded items from all encoders must have the same '
                        'keys if ``multilevel==True``.')

                encoded = {
                    k: [_encoded[k] for _encoded in encoded_list]
                    for k in keys
                }

            else:
                # For combined encoding, the encoded items from different
                # encoders should have no overlapping items, except for
                # `keypoint_weights`. If multiple `keypoint_weights` are given,
                # they will be multiplied as the final `keypoint_weights`.

                encoded = dict()
                keypoint_weights = []

                for _encoded in encoded_list:
                    for key, value in _encoded.items():
                        if key == 'keypoint_weights':
                            keypoint_weights.append(value)
                        elif key not in encoded:
                            encoded[key] = value
                        else:
                            raise ValueError(
                                f'Overlapping item "{key}" from multiple '
                                'encoders, which is not supported when '
                                '``multilevel==False``')

                if keypoint_weights:
                    encoded['keypoint_weights'] = keypoint_weights

            if _field_mapping_table:
                encoded['field_mapping_table'] = _field_mapping_table
            if _instance_mapping_table:
                encoded['instance_mapping_table'] = _instance_mapping_table
            if _label_mapping_table:
                encoded['label_mapping_table'] = _label_mapping_table

        if self.use_dataset_keypoint_weights and 'keypoint_weights' in encoded:
            if isinstance(encoded['keypoint_weights'], list):
                for w in encoded['keypoint_weights']:
                    w = w * results['dataset_keypoint_weights']
            else:
                encoded['keypoint_weights'] = encoded[
                    'keypoint_weights'] * results['dataset_keypoint_weights']

        results.update(encoded)

        return results

    def __repr__(self) -> str:
        """print the basic information of the transform.

        Returns:
            str: Formatted string.
        """
        repr_str = self.__class__.__name__
        repr_str += (f'(encoder={str(self.encoder_cfg)}, ')
        repr_str += ('use_dataset_keypoint_weights='
                     f'{self.use_dataset_keypoint_weights})')
        return repr_str


@TRANSFORMS.register_module()
class YOLOXHSVRandomAug(BaseTransform):
    """Apply HSV augmentation to image sequentially. It is referenced from
    https://github.com/Megvii-
    BaseDetection/YOLOX/blob/main/yolox/data/data_augment.py#L21.

    Required Keys:

    - img

    Modified Keys:

    - img

    Args:
        hue_delta (int): delta of hue. Defaults to 5.
        saturation_delta (int): delta of saturation. Defaults to 30.
        value_delta (int): delat of value. Defaults to 30.
    """

    def __init__(self,
                 hue_delta: int = 5,
                 saturation_delta: int = 30,
                 value_delta: int = 30) -> None:
        self.hue_delta = hue_delta
        self.saturation_delta = saturation_delta
        self.value_delta = value_delta

    @cache_randomness
    def _get_hsv_gains(self):
        hsv_gains = np.random.uniform(-1, 1, 3) * [
            self.hue_delta, self.saturation_delta, self.value_delta
        ]
        # random selection of h, s, v
        hsv_gains *= np.random.randint(0, 2, 3)
        # prevent overflow
        hsv_gains = hsv_gains.astype(np.int16)
        return hsv_gains

    def transform(self, results: dict) -> dict:
        img = results['img']
        hsv_gains = self._get_hsv_gains()
        img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.int16)

        img_hsv[..., 0] = (img_hsv[..., 0] + hsv_gains[0]) % 180
        img_hsv[..., 1] = np.clip(img_hsv[..., 1] + hsv_gains[1], 0, 255)
        img_hsv[..., 2] = np.clip(img_hsv[..., 2] + hsv_gains[2], 0, 255)
        cv2.cvtColor(img_hsv.astype(img.dtype), cv2.COLOR_HSV2BGR, dst=img)

        results['img'] = img
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(hue_delta={self.hue_delta}, '
        repr_str += f'saturation_delta={self.saturation_delta}, '
        repr_str += f'value_delta={self.value_delta})'
        return repr_str


@TRANSFORMS.register_module()
class FilterAnnotations(BaseTransform):
    """Eliminate undesirable annotations based on specific conditions.

    This class is designed to sift through annotations by examining multiple
    factors such as the size of the bounding box, the visibility of keypoints,
    and the overall area. Users can fine-tune the criteria to filter out
    instances that have excessively small bounding boxes, insufficient area,
    or an inadequate number of visible keypoints.

    Required Keys:

    - bbox (np.ndarray) (optional)
    - area (np.int64) (optional)
    - keypoints_visible (np.ndarray) (optional)

    Modified Keys:

    - bbox (optional)
    - bbox_score (optional)
    - category_id (optional)
    - keypoints (optional)
    - keypoints_visible (optional)
    - area (optional)

    Args:
        min_gt_bbox_wh (tuple[float]): Minimum width and height of ground
            truth boxes. Default: (1., 1.)
        min_gt_area (int): Minimum foreground area of instances.
            Default: 1
        min_kpt_vis (int): Minimum number of visible keypoints. Default: 1
        by_box (bool): Filter instances with bounding boxes not meeting the
            min_gt_bbox_wh threshold. Default: False
        by_area (bool): Filter instances with area less than min_gt_area
            threshold. Default: False
        by_kpt (bool): Filter instances with keypoints_visible not meeting the
            min_kpt_vis threshold. Default: True
        keep_empty (bool): Whether to return None when it
            becomes an empty bbox after filtering. Defaults to True.
    """

    def __init__(self,
                 min_gt_bbox_wh: Tuple[int, int] = (1, 1),
                 min_gt_area: int = 1,
                 min_kpt_vis: int = 1,
                 by_box: bool = False,
                 by_area: bool = False,
                 by_kpt: bool = True,
                 keep_empty: bool = True) -> None:

        assert by_box or by_kpt or by_area
        self.min_gt_bbox_wh = min_gt_bbox_wh
        self.min_gt_area = min_gt_area
        self.min_kpt_vis = min_kpt_vis
        self.by_box = by_box
        self.by_area = by_area
        self.by_kpt = by_kpt
        self.keep_empty = keep_empty

    def transform(self, results: dict) -> Union[dict, None]:
        """Transform function to filter annotations.

        Args:
            results (dict): Result dict.

        Returns:
            dict: Updated result dict.
        """
        assert 'keypoints' in results
        kpts = results['keypoints']
        if kpts.shape[0] == 0:
            return results

        tests = []
        if self.by_box and 'bbox' in results:
            bbox = results['bbox']
            tests.append(
                ((bbox[..., 2] - bbox[..., 0] > self.min_gt_bbox_wh[0]) &
                 (bbox[..., 3] - bbox[..., 1] > self.min_gt_bbox_wh[1])))
        if self.by_area and 'area' in results:
            area = results['area']
            tests.append(area >= self.min_gt_area)
        if self.by_kpt:
            kpts_vis = results['keypoints_visible']
            if kpts_vis.ndim == 3:
                kpts_vis = kpts_vis[..., 0]
            tests.append(kpts_vis.sum(axis=1) >= self.min_kpt_vis)

        keep = tests[0]
        for t in tests[1:]:
            keep = keep & t

        if not keep.any():
            if self.keep_empty:
                return None

        keys = ('bbox', 'bbox_score', 'category_id', 'keypoints',
                'keypoints_visible', 'area')
        for key in keys:
            if key in results:
                results[key] = results[key][keep]

        return results

    def __repr__(self):
        return (f'{self.__class__.__name__}('
                f'min_gt_bbox_wh={self.min_gt_bbox_wh}, '
                f'min_gt_area={self.min_gt_area}, '
                f'min_kpt_vis={self.min_kpt_vis}, '
                f'by_box={self.by_box}, '
                f'by_area={self.by_area}, '
                f'by_kpt={self.by_kpt}, '
                f'keep_empty={self.keep_empty})')


def compute_paddings(bbox, bbox_s, kpts):
    """Compute the padding of the bbox to fit the keypoints."""
    bbox = np.array(bbox).flatten()
    bbox_s = np.array(bbox_s).flatten()
    if kpts.size % 2 == 0:
        kpts = kpts.reshape(-1, 2)
    else:
        kpts = kpts.reshape(-1, 3)
    
    x0, y0, x1, y1 = bbox
    x_bbox_distances = np.max(np.stack([
        np.clip(x0 - kpts[:, 0], a_min=0, a_max=None),
        np.clip(kpts[:, 0] - x1, a_min=0, a_max=None),
    ]), axis=0)
    y_bbox_distances = np.max(np.stack([
        np.clip(y0 - kpts[:, 1], a_min=0, a_max=None),
        np.clip(kpts[:, 1] - y1, a_min=0, a_max=None),
    ]), axis=0)

    padding_x = 2 * x_bbox_distances / bbox_s[0]
    padding_y = 2 * y_bbox_distances / bbox_s[1]
    padding = 1 + np.maximum(padding_x, padding_y)

    padding = np.maximum(x_bbox_distances, y_bbox_distances)

    return padding.flatten()