deeplabcut.pose_estimation_pytorch.data.transforms

Classes:

Name	Description
CoarseDropout
Grayscale
HFlip	Horizontal Flip which swaps symmetric keypoints.
KeepAspectRatioResize	Resizes images while preserving their aspect ratio.
KeypointAwareCrop	Random crop for an image around keypoints.
RandomBBoxTransform	Random jittering for bounding boxes for top-down pose estimation models.

Functions:

Name	Description
build_auto_padding	Create an albumentations PadIfNeeded transform from a config.

CoarseDropout

Bases: CoarseDropout

Source code in deeplabcut/pose_estimation_pytorch/data/transforms.py

class CoarseDropout(A.CoarseDropout):
    def __init__(
        self,
        max_holes: int = 8,
        max_height: int | float = 8,
        max_width: int | float = 8,
        min_holes: int | None = None,
        min_height: int | float | None = None,
        min_width: int | float | None = None,
        fill_value: int = 0,
        mask_fill_value: int | None = None,
        always_apply: bool = False,
        p: float = 0.5,
    ):
        super().__init__(
            max_holes,
            max_height,
            max_width,
            min_holes,
            min_height,
            min_width,
            fill_value,
            mask_fill_value,
            always_apply,
            p,
        )

    def apply_to_bboxes(self, bboxes: Sequence[float], **params) -> list[float]:
        return list(bboxes)

    def apply_to_keypoints(
        self,
        keypoints: Sequence[float],
        holes: Iterable[tuple[int, int, int, int]] = (),
        **params,
    ) -> list[float]:
        new_keypoints = []
        for kp in keypoints:
            in_hole = False
            for hole in holes:
                if self._keypoint_in_hole(kp, hole):
                    in_hole = True
                    break
            if in_hole:
                kp = list(kp)
                kp[:2] = np.nan, np.nan
                kp = tuple(kp)
            new_keypoints.append(kp)
        return new_keypoints

    def _keypoint_in_hole(self, keypoint, hole: tuple[int, int, int, int]) -> bool:
        """Reimplemented from Albumentations as was removed in v1.4.0."""
        x1, y1, x2, y2 = hole
        x, y = keypoint[:2]
        return x1 <= x < x2 and y1 <= y < y2

Grayscale

Bases: ToGray

Methods:

Name	Description
__init__	Args:

Source code in deeplabcut/pose_estimation_pytorch/data/transforms.py

class Grayscale(A.ToGray):
    def __init__(
        self,
        alpha: float | int | tuple[float, float] = 1.0,
        always_apply: bool = False,
        p: float = 0.5,
    ):
        """
        Args:
            alpha: int, float or tuple of floats, optional
            The alpha value of the new colorspace when overlaid over the
            old one. A value close to 1.0 means that mostly the new
            colorspace is visible. A value close to 0.0 means that mostly the
            old image is visible.

            * If a float, exactly that value will be used.
            * If a tuple ``(a, b)``, a random value from the range
              ``a <= x <= b`` will be sampled per image.
        """
        super().__init__(always_apply, p)
        if isinstance(alpha, (float, int)):
            self._alpha = self._validate_alpha(alpha)
        elif isinstance(alpha, tuple):
            if len(alpha) != 2:
                raise ValueError("`alpha` must be a tuple of two numbers.")
            self._alpha = tuple([self._validate_alpha(val) for val in alpha])
        else:
            raise ValueError("")

    @staticmethod
    def _validate_alpha(val: float) -> float:
        if not 0.0 <= val <= 1.0:
            warnings.warn("`alpha` will be clipped to the interval [0.0, 1.0].", stacklevel=2)
        return min(1.0, max(0.0, val))

    @property
    def alpha(self) -> float:
        if isinstance(self._alpha, float):
            return self._alpha
        return np.random.uniform(*self._alpha)

    def apply(self, img: NDArray, **params) -> NDArray:
        img_gray = super().apply(img, **params)
        alpha = self.alpha
        img_blend = img * (1 - alpha) + img_gray * alpha
        return img_blend.astype(img.dtype)

__init__

__init__(alpha: float | int | tuple[float, float] = 1.0, always_apply: bool = False, p: float = 0.5)

Parameters:

Name	Type	Description	Default
alpha	float | int | tuple[float, float]	int, float or tuple of floats, optional	1.0

Source code in deeplabcut/pose_estimation_pytorch/data/transforms.py

def __init__(
    self,
    alpha: float | int | tuple[float, float] = 1.0,
    always_apply: bool = False,
    p: float = 0.5,
):
    """
    Args:
        alpha: int, float or tuple of floats, optional
        The alpha value of the new colorspace when overlaid over the
        old one. A value close to 1.0 means that mostly the new
        colorspace is visible. A value close to 0.0 means that mostly the
        old image is visible.

        * If a float, exactly that value will be used.
        * If a tuple ``(a, b)``, a random value from the range
          ``a <= x <= b`` will be sampled per image.
    """
    super().__init__(always_apply, p)
    if isinstance(alpha, (float, int)):
        self._alpha = self._validate_alpha(alpha)
    elif isinstance(alpha, tuple):
        if len(alpha) != 2:
            raise ValueError("`alpha` must be a tuple of two numbers.")
        self._alpha = tuple([self._validate_alpha(val) for val in alpha])
    else:
        raise ValueError("")

HFlip

Bases: HorizontalFlip

Horizontal Flip which swaps symmetric keypoints.

Source code in deeplabcut/pose_estimation_pytorch/data/transforms.py

class HFlip(A.HorizontalFlip):
    """Horizontal Flip which swaps symmetric keypoints."""

    def __init__(self, symmetries: list[tuple[int, int]], *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._symmetries = {}
        for i, j in symmetries:
            self._symmetries[i] = j
            self._symmetries[j] = i

    def apply_to_keypoints(self, keypoints, **params):
        swapped_keypoints = [keypoints[self._symmetries.get(kpt_idx, kpt_idx)] for kpt_idx in range(len(keypoints))]
        return super().apply_to_keypoints(swapped_keypoints, **params)

KeepAspectRatioResize

Bases: DualTransform

Resizes images while preserving their aspect ratio.

In 'pad' mode, the image will be rescaled to the largest possible size such that it can be padded to the correct size (with PadIfNeeded). So we'll have: output_width <= width, output_height <= height

In 'crop' mode, the image will be rescaled to the smallest possible size such that it can be cropped to the correct size (with any random crop you want), so: output_width >= width, output_height >= height

Source code in deeplabcut/pose_estimation_pytorch/data/transforms.py

class KeepAspectRatioResize(A.DualTransform):
    """Resizes images while preserving their aspect ratio.

    In 'pad' mode, the image will be rescaled to the largest possible size such that it
    can be padded to the correct size (with PadIfNeeded). So we'll have: output_width <=
    width, output_height <= height

    In 'crop' mode, the image will be rescaled to the smallest possible size such that
    it can be cropped to the correct size (with any random crop you want), so:
    output_width >= width, output_height >= height
    """

    def __init__(
        self,
        width: int,
        height: int,
        mode: str = "pad",
        interpolation: Any = cv2.INTER_LINEAR,
        p: float = 1.0,
        always_apply: bool = True,
    ) -> None:
        super().__init__(always_apply=always_apply, p=p)
        self.height = height
        self.width = width
        self.mode = mode
        self.interpolation = interpolation

    def apply(self, img, scale=0, interpolation=cv2.INTER_LINEAR, **params):
        return A.scale(img, scale, interpolation)

    def apply_to_bbox(self, bbox, **params):
        # Bounding box coordinates are scale invariant
        return bbox

    def apply_to_keypoint(self, keypoint, scale=0, **params):
        keypoint = A.keypoint_scale(keypoint, scale, scale)
        return keypoint

    @property
    def targets_as_params(self) -> list[str]:
        return ["image"]

    def get_params_dependent_on_targets(self, params: dict[str, Any]) -> dict[str, Any]:
        h, w, _ = params["image"].shape
        if self.mode == "pad":
            scale = min(self.height / h, self.width / w)
        else:
            scale = max(self.height / h, self.width / w)

        return {"scale": scale}

    def get_transform_init_args_names(self):
        return "height", "width", "mode", "interpolation"

KeypointAwareCrop

Bases: RandomCrop

Random crop for an image around keypoints.

Parameters:

Name	Type	Description	Default
width	int	Crop images down to this maximum width.	required
height	int	Crop images down to this maximum height.	required
max_shift	float	Maximum allowed shift of the cropping center position as a fraction of the crop size.	0.4
crop_sampling	str	Crop centers sampling method. Must be either: "uniform" (randomly over the image), "keypoints" (randomly over the annotated keypoints), "density" (weighing preferentially dense regions of keypoints), "hybrid" (alternating randomly between "uniform" and "density").	'hybrid'

Source code in deeplabcut/pose_estimation_pytorch/data/transforms.py

class KeypointAwareCrop(A.RandomCrop):
    """Random crop for an image around keypoints.

    Args:
        width: Crop images down to this maximum width.
        height: Crop images down to this maximum height.
        max_shift: Maximum allowed shift of the cropping center position
            as a fraction of the crop size.
        crop_sampling: Crop centers sampling method. Must be either:
            "uniform" (randomly over the image),
            "keypoints" (randomly over the annotated keypoints),
            "density" (weighing preferentially dense regions of keypoints),
            "hybrid" (alternating randomly between "uniform" and "density").
    """

    def __init__(
        self,
        width: int,
        height: int,
        max_shift: float = 0.4,
        crop_sampling: str = "hybrid",
    ):
        super().__init__(height, width, always_apply=True)
        # Clamp to 40% of crop size to ensure that at least
        # the center keypoint remains visible after the offset is applied.
        self.max_shift = max(0.0, min(max_shift, 0.4))
        if crop_sampling not in ("uniform", "keypoints", "density", "hybrid"):
            raise ValueError(
                f"Invalid sampling {crop_sampling}. Must be either 'uniform', 'keypoints', 'density', or 'hybrid."
            )
        self.crop_sampling = crop_sampling

    @staticmethod
    def calc_n_neighbors(xy: NDArray, radius: float) -> NDArray:
        d = pdist(xy, "sqeuclidean")
        mat = squareform(d <= radius * radius, checks=False)
        return np.sum(mat, axis=0)

    @property
    def targets_as_params(self) -> list[str]:
        return ["image", "keypoints"]

    def get_params_dependent_on_targets(self, params: dict[str, Any]) -> dict[str, Any]:
        img = params["image"]
        kpts = params["keypoints"]
        shift_factors = np.random.random(2)
        shift = self.max_shift * shift_factors * np.array([self.width, self.height])
        sampling = self.crop_sampling
        if self.crop_sampling == "hybrid":
            sampling = np.random.choice(["uniform", "density"])
        if len(kpts) == 0:
            sampling = "uniform"
        if sampling == "uniform":
            center = np.random.random(2)
        else:
            h, w = img.shape[:2]
            kpts = np.array([[k[0], k[1]] for k in kpts])
            kpts = kpts[~np.isnan(kpts).all(axis=1)]
            n_kpts = kpts.shape[0]
            inds = np.arange(n_kpts)
            if sampling == "density":
                # Points located close to one another are sampled preferentially
                # in order to augment crowded regions.
                radius = 0.1 * min(h, w)
                n_neighbors = self.calc_n_neighbors(kpts, radius)
                # Include keypoints in the count to avoid null probabilities
                n_neighbors += 1
                p = n_neighbors / n_neighbors.sum()
            else:
                p = np.ones_like(inds) / n_kpts
            center = kpts[np.random.choice(inds, p=p)]
            # Shift the crop center in both dimensions by random amounts
            # and normalize to the original image dimensions.
            center = (center + shift) / [w, h]
            center = np.clip(center, 0, np.nextafter(1, 0))  # Clip to 1 exclusive
        return {"h_start": center[1], "w_start": center[0]}

    def apply_to_keypoints(
        self,
        keypoints,
        **params,
    ) -> list[tuple[float]]:
        keypoints = super().apply_to_keypoints(keypoints, **params)
        new_keypoints = []
        for kp in keypoints:
            x, y = kp[:2]
            if not (0 <= x < self.width and 0 <= y < self.height):
                kp = list(kp)
                kp[:2] = np.nan, np.nan
                kp = tuple(kp)
            new_keypoints.append(kp)
        return new_keypoints

    def get_transform_init_args_names(self) -> tuple[str, ...]:
        return "width", "height", "max_shift", "crop_sampling"

RandomBBoxTransform

Bases: DualTransform

Random jittering for bounding boxes for top-down pose estimation models.

Implementation based on the mmpose RandomBBoxTransform. For more information, see https://github.com/open-mmlab/mmpose.

Source code in deeplabcut/pose_estimation_pytorch/data/transforms.py

class RandomBBoxTransform(A.DualTransform):
    """Random jittering for bounding boxes for top-down pose estimation models.

    Implementation based on the mmpose `RandomBBoxTransform`. For more information,
    see <https://github.com/open-mmlab/mmpose>.
    """

    def __init__(
        self,
        shift_factor: float = 0.1,
        shift_prob: float = 0.25,
        scale_factor: tuple[float, float] = (0.5, 1.5),
        scale_prob: float = 1.0,
        sampling: str = "truncnorm",
        p: float = 1.0,
    ):
        super().__init__(p=p)
        self.shift_factor = shift_factor
        self.shift_prob = shift_prob
        self.scale_factor = scale_factor
        self.scale_prob = scale_prob
        self.sampling = sampling

    def apply(self, img: np.ndarray, **params) -> np.ndarray:
        return img

    def apply_to_keypoints(self, keypoints: np.ndarray, **params) -> np.ndarray:
        return keypoints

    def apply_to_bboxes(self, bboxes, **params):
        if len(bboxes) == 0:
            return bboxes

        # Albumentations provides bounding boxes in normalized xyxy format internally
        bboxes_xyxy = np.asarray(bboxes)
        bboxes_extra = None
        if bboxes_xyxy.shape[1] > 4:
            # can't take from array - may have different dtype
            bboxes_extra = [bbox[4:] for bbox in bboxes]
            bboxes_xyxy = bboxes_xyxy[:, :4]

        # sample parameters
        bboxes_to_scale = np.random.random(len(bboxes)) < self.scale_prob
        num_bboxes_to_scale = np.sum(bboxes_to_scale).item()
        scale_factors = np.ones((len(bboxes), 2))
        if num_bboxes_to_scale > 0:
            scale_factors[bboxes_to_scale] = self._sample(
                (num_bboxes_to_scale, 2),
                low=self.scale_factor[0],
                high=self.scale_factor[1],
            )

        bboxes_to_shift = np.random.random(len(bboxes)) < self.shift_prob
        num_bboxes_to_shift = np.sum(bboxes_to_shift).item()
        shift_factors = np.zeros((len(bboxes), 2))
        if num_bboxes_to_shift > 0:
            shift_factors[bboxes_to_shift] = self._sample(
                (num_bboxes_to_shift, 2),
                low=-self.shift_factor,
                high=self.shift_factor,
            )

        bbox_wh = bboxes_xyxy[:, 2:] - bboxes_xyxy[:, :2]
        bbox_cxcy = bboxes_xyxy[:, :2] + (0.5 * bbox_wh)

        # scale + shift bounding boxes
        bbox_cxcy = bbox_cxcy + (shift_factors * bbox_wh)
        bbox_wh = bbox_wh * scale_factors

        # convert to xyxy, clip so all bounding boxes are in the image
        bbox_half_wh = 0.5 * bbox_wh
        bbox_xyxy = np.empty((len(bboxes), 4))
        bbox_xyxy[:, :2] = bbox_cxcy - bbox_half_wh
        bbox_xyxy[:, 2:] = bbox_cxcy + bbox_half_wh
        bbox_xyxy = np.clip(bbox_xyxy, 0, 1)

        # add the extra information back; tuples for albumentations<=1.4.3
        bboxes_out = [tuple(bbox) for bbox in bbox_xyxy]
        if bboxes_extra is not None:
            bboxes_out = [bbox + extra for bbox, extra in zip(bboxes_out, bboxes_extra, strict=False)]
        return bboxes_out

    def get_transform_init_args_names(self):
        return "shift_factor", "shift_prob", "scale_factor", "scale_prob", "sampling"

    def _sample(
        self,
        size: tuple[int, ...],
        low: float = -1.0,
        high: float = 1.0,
    ) -> np.ndarray:
        if self.sampling == "truncnorm":
            return truncnorm.rvs(low, high, size=size).astype(np.float32)
        elif self.sampling == "uniform":
            delta = high - low
            return low + (delta * np.random.random(size))

        raise ValueError(f"Unknown sampling: {self.sampling}")

build_auto_padding

build_auto_padding(
    min_height: int | None = None,
    min_width: int | None = None,
    pad_height_divisor: int | None = 1,
    pad_width_divisor: int | None = 1,
    position: str = "random",
    border_mode: str = "reflect_101",
    border_value: float | None = None,
    border_mask_value: float | None = None,
) -> A.PadIfNeeded

Create an albumentations PadIfNeeded transform from a config.

Parameters:

Name	Type	Description	Default
min_height	int | None	the minimum height of the image	None
min_width	int | None	the minimum width of the image	None
pad_height_divisor	int | None	if not None, ensures height is dividable by value of this argument	1
pad_width_divisor	int | None	if not None, ensures width is dividable by value of this argument	1
position	str	position of the image, one of the possible PadIfNeeded	'random'
border_mode	str	'constant' or 'reflect_101' (see cv2.BORDER modes)	'reflect_101'
border_value	float | None	padding value if border_mode is 'constant'	None
border_mask_value	float | None	padding value for mask if border_mode is 'constant'	None

Raises:

Type	Description
ValueError	Only one of 'min_height' and 'pad_height_divisor' parameters must be set Only one of 'min_width' and 'pad_width_divisor' parameters must be set

Returns:

Type	Description
PadIfNeeded	the auto-padding transform

Source code in deeplabcut/pose_estimation_pytorch/data/transforms.py

def build_auto_padding(
    min_height: int | None = None,
    min_width: int | None = None,
    pad_height_divisor: int | None = 1,
    pad_width_divisor: int | None = 1,
    position: str = "random",  # TODO: Which default to set?
    border_mode: str = "reflect_101",  # TODO: Which default to set?
    border_value: float | None = None,
    border_mask_value: float | None = None,
) -> A.PadIfNeeded:
    """Create an albumentations PadIfNeeded transform from a config.

    Args:
        min_height: the minimum height of the image
        min_width: the minimum width of the image
        pad_height_divisor: if not None, ensures height is dividable by value of this argument
        pad_width_divisor: if not None, ensures width is dividable by value of this argument
        position: position of the image, one of the possible PadIfNeeded
        border_mode: 'constant' or 'reflect_101' (see cv2.BORDER modes)
        border_value: padding value if border_mode is 'constant'
        border_mask_value: padding value for mask if border_mode is 'constant'

    Raises:
        ValueError:
            Only one of 'min_height' and 'pad_height_divisor' parameters must be set
            Only one of 'min_width' and 'pad_width_divisor' parameters must be set

    Returns:
        the auto-padding transform
    """
    border_modes = {
        "constant": cv2.BORDER_CONSTANT,
        "reflect_101": cv2.BORDER_REFLECT_101,
    }
    if border_mode not in border_modes:
        raise ValueError(
            f"Unknown border mode for auto_padding: {border_mode} (valid values are: {border_modes.keys()})"
        )

    return A.PadIfNeeded(
        min_height=min_height,
        min_width=min_width,
        pad_height_divisor=pad_height_divisor,
        pad_width_divisor=pad_width_divisor,
        position=position,
        border_mode=border_modes[border_mode],
        value=border_value,
        mask_value=border_mask_value,
    )