deeplabcut.pose_estimation_pytorch.data.utils

Functions:

Name	Description
apply_transform	Applies a transformation to the provided image and keypoints.
bbox_from_keypoints	Computes bounding boxes from keypoints.
calc_area_from_keypoints	Calculate the area from keypoints.
calc_bbox_overlap	Calculate the overlap between two bounding boxes.
map_id_to_annotations	Maps image IDs to their corresponding annotation indices.
map_image_path_to_id	Binds the image paths to their respective IDs.
merge_list_of_dicts	Flattens a list of dictionaries into a dictionary with the lists concatenated.
out_of_bounds_keypoints	Computes which visible keypoints are outside an image.
pad_to_length	Pads the first dimension of an array with a given value.
read_image_shape_fast	Blazing fast and does not load the image into memory.
safe_stack	Stacks a list of arrays if there are any, otherwise returns an array of zeros of

apply_transform

apply_transform(
    transform: BaseCompose, image: ndarray, keypoints: ndarray, bboxes: ndarray, class_labels: list[str]
) -> dict[str, np.ndarray]

Applies a transformation to the provided image and keypoints.

Parameters:

Name	Type	Description	Default
transform	BaseCompose	The transformation to apply.	required
image	ndarray	The input image to which the transformation will be applied.	required
keypoints	ndarray	List of keypoints to be transformed along with the image. Each keypoint is expected to be a tuple or list with at least three values, where the third value indicates the class label index.	required
bboxes	ndarray	List of bounding boxes to be transformed along with the image.	required
class_labels	list[str]	List of class labels corresponding to the keypoints.	required

Returns:

Name	Type	Description
transformed	dict[str, ndarray]	A dictionary containing the transformed image and keypoints.

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

def apply_transform(
    transform: A.BaseCompose,
    image: np.ndarray,
    keypoints: np.ndarray,
    bboxes: np.ndarray,
    class_labels: list[str],
) -> dict[str, np.ndarray]:
    """Applies a transformation to the provided image and keypoints.

    Args:
        transform: The transformation to apply.
        image: The input image to which the transformation will be applied.
        keypoints: List of keypoints to be transformed along with the image. Each keypoint
            is expected to be a tuple or list with at least three values,
            where the third value indicates the class label index.
        bboxes: List of bounding boxes to be transformed along with the image.
        class_labels: List of class labels corresponding to the keypoints.

    Returns:
        transformed: A dictionary containing the transformed image and keypoints.
    """

    if transform:
        oob_mask = out_of_bounds_keypoints(keypoints, image.shape)
        transformed = _apply_transform(transform, image, keypoints, bboxes, class_labels)

        transformed["keypoints"] = np.array(transformed["keypoints"])

        # out-of-bound keypoints have visibility flag 0. But we don't touch coordinates
        if np.sum(oob_mask) > 0:
            transformed["keypoints"][oob_mask, 2] = 0.0

        out_shape = transformed["image"].shape
        if len(transformed["keypoints"]) > 0:
            oob_mask = out_of_bounds_keypoints(transformed["keypoints"], out_shape)
            # out-of-bound keypoints have visibility flag 0. Don't touch coordinates
            if np.sum(oob_mask) > 0:
                transformed["keypoints"][oob_mask, 2] = 0.0

        # TODO: Check that the transformed bboxes are still within the image
        if len(transformed["bboxes"]) > 0:
            transformed["bboxes"] = np.array(transformed["bboxes"])
        else:
            transformed["bboxes"] = np.zeros(shape=(0, 4))

    else:
        transformed = {"keypoints": keypoints, "image": image}

    # do we ever need to do this if we had check_keypoints_within_bounds above?
    # np.nan_to_num(transformed["keypoints"], copy=False, nan=-1)
    return transformed

bbox_from_keypoints

bbox_from_keypoints(keypoints: ndarray, image_h: int, image_w: int, margin: int) -> np.ndarray

Computes bounding boxes from keypoints.

Parameters:

Name	Type	Description	Default
keypoints	ndarray	(..., num_keypoints, xy) the keypoints from which to get bboxes	required
image_h	int	the height of the image	required
image_w	int	the width of the image	required
margin	int	the bounding box margin	required

Returns:

Type	Description
ndarray	the bounding boxes for the keypoints, of shape (..., 4) in the xywh format

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

def bbox_from_keypoints(
    keypoints: np.ndarray,
    image_h: int,
    image_w: int,
    margin: int,
) -> np.ndarray:
    """Computes bounding boxes from keypoints.

    Args:
        keypoints: (..., num_keypoints, xy) the keypoints from which to get bboxes
        image_h: the height of the image
        image_w: the width of the image
        margin: the bounding box margin

    Returns:
        the bounding boxes for the keypoints, of shape (..., 4) in the xywh format
    """
    squeeze = False

    # we do not estimate bbox on keypoints that have 0 or -1 flag
    keypoints = np.copy(keypoints)
    keypoints[keypoints[..., -1] <= 0] = np.nan

    if len(keypoints.shape) == 2:
        squeeze = True
        keypoints = np.expand_dims(keypoints, axis=0)

    bboxes = np.full((keypoints.shape[0], 4), np.nan)
    with warnings.catch_warnings():  # silence warnings when all pose confidence levels are <= 0
        warnings.simplefilter("ignore", category=RuntimeWarning)
        bboxes[:, :2] = np.nanmin(keypoints[..., :2], axis=1) - margin  # X1, Y1
        bboxes[:, 2:4] = np.nanmax(keypoints[..., :2], axis=1) + margin  # X2, Y2

    # can have NaNs if some individuals have no visible keypoints
    bboxes = np.nan_to_num(bboxes, nan=0)

    bboxes = np.clip(
        bboxes,
        a_min=[0, 0, 0, 0],
        a_max=[image_w, image_h, image_w, image_h],
    )
    bboxes[..., 2] = bboxes[..., 2] - bboxes[..., 0]  # to width
    bboxes[..., 3] = bboxes[..., 3] - bboxes[..., 1]  # to height
    if squeeze:
        return bboxes[0]

    return bboxes

calc_area_from_keypoints

calc_area_from_keypoints(keypoints: ndarray) -> np.ndarray

Calculate the area from keypoints.

in the pups benchmark, there are 5 keypoints perfectly aligned so

the area is 0. How do we deal with that? Makes more sense to compute the area from the bboxes (they are padded) Below is a temporary fix, which sets a min height and width to 5 Suggestion: compute min height/width using labeled data

Parameters:

Name	Type	Description	Default
keypoints	ndarray	array of keypoints	required

Returns:

Type	Description
ndarray	np.ndarray: array containing the computed areas based on the keypoints

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

def calc_area_from_keypoints(keypoints: np.ndarray) -> np.ndarray:
    """Calculate the area from keypoints.

    TODO: in the pups benchmark, there are 5 keypoints perfectly aligned so
     the area is 0.
     How do we deal with that?
     Makes more sense to compute the area from the bboxes (they are padded)
     Below is a temporary fix, which sets a min height and width to 5
     Suggestion: compute min height/width using labeled data

    Args:
        keypoints (np.ndarray): array of keypoints

    Returns:
        np.ndarray: array containing the computed areas based on the keypoints
    """
    w = np.maximum(keypoints[:, :, 0].max(axis=1) - keypoints[:, :, 0].min(axis=1), 1)
    h = np.maximum(keypoints[:, :, 1].max(axis=1) - keypoints[:, :, 1].min(axis=1), 1)
    return w * h

calc_bbox_overlap

calc_bbox_overlap(bbox1: ndarray, bbox2: ndarray) -> np.ndarray

Calculate the overlap between two bounding boxes.

Parameters:

Name	Type	Description	Default
bbox1	ndarray	the first bounding box in the format (x, y, w, h)	required
bbox2	ndarray	the second bounding box in the format (x, y, w, h)	required

Returns:

Type	Description
ndarray	The overlap between

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

def calc_bbox_overlap(bbox1: np.ndarray, bbox2: np.ndarray) -> np.ndarray:
    """Calculate the overlap between two bounding boxes.

    Args:
        bbox1: the first bounding box in the format (x, y, w, h)
        bbox2: the second bounding box in the format (x, y, w, h)

    Returns:
        The overlap between
    """
    x1, y1, w1, h1 = bbox1
    x2, y2, w2, h2 = bbox2

    x1_max = x1 + w1
    y1_max = y1 + h1
    x2_max = x2 + w2
    y2_max = y2 + h2

    x_overlap = max(0, min(x1_max, x2_max) - max(x1, x2))
    y_overlap = max(0, min(y1_max, y2_max) - max(y1, y2))

    intersection = x_overlap * y_overlap
    union = w1 * h1 + w2 * h2 - intersection

    return intersection / union

map_id_to_annotations

map_id_to_annotations(annotations: list[dict]) -> dict[int, list[int]]

Maps image IDs to their corresponding annotation indices.

Parameters:

Name	Type	Description	Default
annotations	list[dict]	List of dictionaries containing annotation data. Each dictionary should have 'image_id' key.	required

Returns:

Type	Description
dict[int, list[int]]	A dictionary mapping image IDs to lists of corresponding annotation indices.

Examples:

annotations = [{"image_id": 1, ...}, ...]

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

def map_id_to_annotations(annotations: list[dict]) -> dict[int, list[int]]:
    """Maps image IDs to their corresponding annotation indices.

    Args:
        annotations: List of dictionaries containing annotation data. Each dictionary
            should have 'image_id' key.

    Returns:
        A dictionary mapping image IDs to lists of corresponding annotation indices.

    Examples:
        annotations = [{"image_id": 1, ...}, ...]
    """

    annotation_idx_map = defaultdict(list)
    for idx, annotation in enumerate(annotations):
        annotation_idx_map[annotation["image_id"]].append(idx)

    return annotation_idx_map

map_image_path_to_id

map_image_path_to_id(images: list[dict]) -> dict[str, int]

Binds the image paths to their respective IDs.

Parameters:

Name	Type	Description	Default
images	list[dict]	List of dictionaries containing image data in COCO-like format. Each dictionary should have 'file_name' and 'id' keys.	required

Returns:

Type	Description
dict[str, int]	A dictionary mapping image paths to their respective IDs.

Examples:

images = [{"file_name": "path/to/image1.jpg", "id": 1}, ...]

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

def map_image_path_to_id(images: list[dict]) -> dict[str, int]:
    """Binds the image paths to their respective IDs.

    Args:
        images: List of dictionaries containing image data in COCO-like format.
            Each dictionary should have 'file_name' and 'id' keys.

    Returns:
        A dictionary mapping image paths to their respective IDs.

    Examples:
        images = [{"file_name": "path/to/image1.jpg", "id": 1}, ...]
    """

    return {image["file_name"]: image["id"] for image in images}

merge_list_of_dicts

merge_list_of_dicts(list_of_dicts: list[dict], keys_to_include: list[str]) -> dict[str, list]

Flattens a list of dictionaries into a dictionary with the lists concatenated.

Parameters:

Name	Type	Description	Default
list_of_dicts	list[dict]	the dictionaries to merge	required
keys_to_include	list[str]	the keys to include in the new dictionary	required

Returns:

Type	Description
dict[str, list]	the merged dictionary

Examples:

input: list_of_dicts: [{"id": 0, "num": 1}, {"id": 1, "num": 10}] keys_to_include: ["id", "num"] output:

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

def merge_list_of_dicts(list_of_dicts: list[dict], keys_to_include: list[str]) -> dict[str, list]:
    """Flattens a list of dictionaries into a dictionary with the lists concatenated.

    Args:
        list_of_dicts: the dictionaries to merge
        keys_to_include: the keys to include in the new dictionary

    Returns:
        the merged dictionary

    Examples:
        input:
            list_of_dicts: [{"id": 0, "num": 1}, {"id": 1, "num": 10}]
            keys_to_include: ["id", "num"]
        output:
            {"id": [0, 1], "num": [1, 10]}
    """
    return reduce(
        lambda acc, d: {key: acc.get(key, []) + [value] for key, value in d.items() if key in keys_to_include},
        list_of_dicts,
        defaultdict(list),
    )

out_of_bounds_keypoints

out_of_bounds_keypoints(keypoints: ndarray, shape: tuple) -> np.ndarray

Computes which visible keypoints are outside an image.

Parameters:

Name	Type	Description	Default
keypoints	ndarray	A (N, 3) shaped array where N is the number of keypoints and each keypoint is represented as (x, y, visibility).	required
shape	tuple	A tuple representing the shape or bounds as (height, width).	required

Returns:

Type	Description
ndarray	A boolean array of shape (N,) where each element corresponds to whether the respective keypoint is visible (visibility > 0) and outside the image bounds. This mask can be used to set the visibility bit to 0 for keypoints that were kicked off an image due to augmentation.

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

def out_of_bounds_keypoints(keypoints: np.ndarray, shape: tuple) -> np.ndarray:
    """Computes which visible keypoints are outside an image.

    Args:
        keypoints: A (N, 3) shaped array where N is the number of keypoints and each
            keypoint is represented as (x, y, visibility).
        shape: A tuple representing the shape or bounds as (height, width).

    Returns:
        A boolean array of shape (N,) where each element corresponds to whether
        the respective keypoint is visible (visibility > 0) and outside the image
        bounds. This mask can be used to set the visibility bit to 0 for keypoints that
        were kicked off an image due to augmentation.
    """
    return (keypoints[..., 2] > 0) & (
        np.isnan(keypoints[..., 0])
        | np.isnan(keypoints[..., 1])
        | (keypoints[..., 0] < 0)
        | (keypoints[..., 0] > shape[1])
        | (keypoints[..., 1] < 0)
        | (keypoints[..., 1] > shape[0])
    )

pad_to_length

pad_to_length(data: array, length: int, value: float) -> np.array

Pads the first dimension of an array with a given value.

Parameters:

Name	Type	Description	Default
data	array	the array to pad, of shape (l, ...), where l <= length	required
length	int	the desired length of the tensor	required
value	float	the value to pad with	required

Returns:

Type	Description
array	the padded array of shape (length, ...)

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

def pad_to_length(data: np.array, length: int, value: float) -> np.array:
    """Pads the first dimension of an array with a given value.

    Args:
        data: the array to pad, of shape (l, ...), where l <= length
        length: the desired length of the tensor
        value: the value to pad with

    Returns:
        the padded array of shape (length, ...)
    """
    pad_length = length - len(data)
    if pad_length == 0:
        return data
    elif pad_length > 0:
        padding = value * np.ones((pad_length, *data.shape[1:]), dtype=data.dtype)
        return np.concatenate([data, padding])

    raise ValueError(f"Cannot pad! data.shape={data.shape} > length={length}")

read_image_shape_fast cached

read_image_shape_fast(path: str | Path) -> tuple[int, int, int]

Blazing fast and does not load the image into memory.

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

@cache
def read_image_shape_fast(path: str | Path) -> tuple[int, int, int]:
    """Blazing fast and does not load the image into memory."""
    with Image.open(path) as img:
        width, height = img.size
        return len(img.getbands()), height, width

safe_stack

safe_stack(data: list[ndarray], default_shape: tuple[int, ...]) -> np.ndarray

Stacks a list of arrays if there are any, otherwise returns an array of zeros of a desired shape.

Parameters:

Name	Type	Description	Default
data	list[ndarray]	the list of arrays to stack	required
default_shape	tuple[int, ...]	the shape of the array to return if the list is empty	required

Returns:

Type	Description
ndarray	the stacked data or empty array

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

def safe_stack(data: list[np.ndarray], default_shape: tuple[int, ...]) -> np.ndarray:
    """Stacks a list of arrays if there are any, otherwise returns an array of zeros of
    a desired shape.

    Args:
        data: the list of arrays to stack
        default_shape: the shape of the array to return if the list is empty

    Returns:
        the stacked data or empty array
    """
    if len(data) == 0:
        return np.zeros(default_shape, dtype=float)

    return np.stack(data, axis=0)