Skip to content

Samples

PoseEstimationSample dataclass

A data class describing a single pose estimation sample that comes from a dataset. It contains both input image and target information to train a pose estimation model.

Parameters:

Name Type Description Default
image Union[np.ndarray, torch.Tensor]

Associated image with a sample. Can be in [H,W,C] or [C,H,W] format

 required
image_layout

Layout of the image (HWC or CHW)

 required
mask Union[np.ndarray, torch.Tensor]

Target mask in [H,W] format

 required
joints np.ndarray

Target joints in [NumInstances, NumJoints, 3] format. Last dimension contains (x,y,visibility) for each joint.

 required
areas Optional[np.ndarray]

(Optional) Numpy array of [N] shape with area of each instance. Note this is not a bbox area, but area of the object itself. One may use a heuristic 0.53 * box area as object area approximation if this is not provided.

 required
bboxes_xywh Optional[np.ndarray]

(Optional) Numpy array of [N,4] shape with bounding box of each instance (XYWH)

 required
additional_samples Optional[List[PoseEstimationSample]]

(Optional) List of additional samples for the same image.

 required
is_crowd Optional[np.ndarray]

(Optional) Numpy array of [N] shape with is_crowd flag for each instance

 required
Source code in V3_4/src/super_gradients/training/samples/pose_estimation_sample.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
@dataclasses.dataclass
class PoseEstimationSample:
    """
    A data class describing a single pose estimation sample that comes from a dataset.
    It contains both input image and target information to train a pose estimation model.

    :param image:              Associated image with a sample. Can be in [H,W,C] or [C,H,W] format
    :param image_layout:       Layout of the image (HWC or CHW)
    :param mask:               Target mask in [H,W] format
    :param joints:             Target joints in [NumInstances, NumJoints, 3] format.
                               Last dimension contains (x,y,visibility) for each joint.
    :param areas:              (Optional) Numpy array of [N] shape with area of each instance.
                               Note this is not a bbox area, but area of the object itself.
                               One may use a heuristic `0.53 * box area` as object area approximation if this is not provided.
    :param bboxes_xywh:        (Optional) Numpy array of [N,4] shape with bounding box of each instance (XYWH)
    :param additional_samples: (Optional) List of additional samples for the same image.
    :param is_crowd:           (Optional) Numpy array of [N] shape with is_crowd flag for each instance
    """

    __slots__ = ["image", "mask", "joints", "areas", "bboxes_xywh", "is_crowd", "additional_samples"]

    image: Union[np.ndarray, torch.Tensor]
    mask: Union[np.ndarray, torch.Tensor]
    joints: np.ndarray
    areas: Optional[np.ndarray]
    bboxes_xywh: Optional[np.ndarray]
    is_crowd: Optional[np.ndarray]
    additional_samples: Optional[List["PoseEstimationSample"]]

    @classmethod
    def compute_area_of_joints_bounding_box(cls, joints) -> np.ndarray:
        """
        Compute area of a bounding box enclosing visible joints for each pose instance.
        :param joints: np.ndarray of [Num Instances, Num Joints, 3] shape (x,y,visibility)
        :return:       np.ndarray of [Num Instances] shape with box area of the visible joints
                       (zero if all joints are not visible or only one joint is visible)
        """
        visible_joints = joints[:, :, 2] > 0
        xmax = np.max(joints[:, :, 0], axis=-1, where=visible_joints, initial=joints[:, :, 0].min())
        xmin = np.min(joints[:, :, 0], axis=-1, where=visible_joints, initial=joints[:, :, 0].max())
        ymax = np.max(joints[:, :, 1], axis=-1, where=visible_joints, initial=joints[:, :, 1].min())
        ymin = np.min(joints[:, :, 1], axis=-1, where=visible_joints, initial=joints[:, :, 1].max())

        w = xmax - xmin
        h = ymax - ymin
        raw_area = w * h
        area = np.clip(raw_area, a_min=0, a_max=None) * (visible_joints.sum(axis=-1, keepdims=False) > 1)
        return area

    def sanitize_sample(self) -> "PoseEstimationSample":
        """
        Apply sanity checks on the pose sample, which includes:
        - Clamp bbox coordinates to ensure they are within image boundaries
        - Update visibility status of keypoints if they are outside of image boundaries
        - Update area if bbox clipping occurs
        This function does not remove instances, but may make them subject for removal instead.
        :return: self
        """
        image_height, image_width, _ = self.image.shape

        # Update joints visibility status
        outside_left = self.joints[:, :, 0] < 0
        outside_top = self.joints[:, :, 1] < 0
        outside_right = self.joints[:, :, 0] >= image_width
        outside_bottom = self.joints[:, :, 1] >= image_height

        outside_image_mask = outside_left | outside_top | outside_right | outside_bottom
        self.joints[outside_image_mask, 2] = 0

        if self.bboxes_xywh is not None:
            # Clamp bboxes to image boundaries
            clamped_boxes = xywh_to_xyxy(self.bboxes_xywh, image_shape=(image_height, image_width))
            clamped_boxes[..., [0, 2]] = np.clip(clamped_boxes[..., [0, 2]], 0, image_width - 1)
            clamped_boxes[..., [1, 3]] = np.clip(clamped_boxes[..., [1, 3]], 0, image_height - 1)
            clamped_boxes = xyxy_to_xywh(clamped_boxes, image_shape=(image_height, image_width))

            # Recompute sample areas if they are present
            if self.areas is not None:
                area_reduction_factor = clamped_boxes[..., 2:4].prod(axis=-1) / (self.bboxes_xywh[..., 2:4].prod(axis=-1) + 1e-6)
                self.areas = self.areas * area_reduction_factor

            self.bboxes_xywh = clamped_boxes
        return self

    def filter_by_mask(self, mask: np.ndarray) -> "PoseEstimationSample":
        """
        Remove pose instances with respect to given mask.

        :remark: This is main method to modify instances of the sample.
        If you are implementing a subclass of PoseEstimationSample and adding extra field associated with each pose
        instance (Let's say you add a distance property for each pose from the camera), then you should override
        this method to do filtering on extra attribute as well.

        :param mask:   A boolean or integer mask of samples to keep for given sample.
        :return:       A pose sample after filtering.
        """
        self.joints = self.joints[mask]
        self.is_crowd = self.is_crowd[mask]
        if self.bboxes_xywh is not None:
            self.bboxes_xywh = self.bboxes_xywh[mask]
        if self.areas is not None:
            self.areas = self.areas[mask]
        return self

    def filter_by_visible_joints(self, min_visible_joints: int) -> "PoseEstimationSample":
        """
        Remove instances from the sample which has less than N visible joints.

        :param min_visible_joints: A minimal number of visible joints a pose has to have in order to be kept.
        :return:                   A pose sample after filtering.
        """
        visible_joints_mask = self.joints[:, :, 2] > 0
        keep_mask: np.ndarray = np.sum(visible_joints_mask, axis=-1) >= min_visible_joints
        return self.filter_by_mask(keep_mask)

    def filter_by_bbox_area(self, min_bbox_area: Union[int, float]) -> "PoseEstimationSample":
        """
        Remove pose instances that has area of the corresponding bounding box less than a certain threshold.

        :param sample:        Instance of PoseEstimationSample to modify. Modification done in-place.
        :param min_bbox_area: Minimal bounding box area of the pose to keep.
        :return:              A pose sample after filtering.
        """
        if self.bboxes_xywh is None:
            area = self.compute_area_of_joints_bounding_box(self.joints)
        else:
            area = self.bboxes_xywh[..., 2:4].prod(axis=-1)

        keep_mask = area >= min_bbox_area
        return self.filter_by_mask(keep_mask)

    def filter_by_pose_area(self, min_instance_area: Union[int, float]) -> "PoseEstimationSample":
        """
        Remove pose instances which area is less than a certain threshold.

        :param sample:            Instance of PoseEstimationSample to modify. Modification done in-place.
        :param min_instance_area: Minimal area of the pose to keep.
        :return:                  A pose sample after filtering.
        """

        if self.areas is not None:
            areas = self.areas
        elif self.bboxes_xywh is not None:
            # 0.53 is a heuristic multiplier from COCO to approximate object area from bbox area
            areas = self.bboxes_xywh[..., 2:4].prod(axis=-1, keepdims=False) * 0.53
        else:
            areas = self.compute_area_of_joints_bounding_box(self.joints)

        keep_mask = areas >= min_instance_area
        return self.filter_by_mask(keep_mask)

compute_area_of_joints_bounding_box(joints) classmethod

Compute area of a bounding box enclosing visible joints for each pose instance.

Parameters:

Name Type Description Default
joints

np.ndarray of [Num Instances, Num Joints, 3] shape (x,y,visibility)

 required

Returns:

Type Description
np.ndarray

np.ndarray of [Num Instances] shape with box area of the visible joints (zero if all joints are not visible or only one joint is visible)
Source code in V3_4/src/super_gradients/training/samples/pose_estimation_sample.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
@classmethod
def compute_area_of_joints_bounding_box(cls, joints) -> np.ndarray:
    """
    Compute area of a bounding box enclosing visible joints for each pose instance.
    :param joints: np.ndarray of [Num Instances, Num Joints, 3] shape (x,y,visibility)
    :return:       np.ndarray of [Num Instances] shape with box area of the visible joints
                   (zero if all joints are not visible or only one joint is visible)
    """
    visible_joints = joints[:, :, 2] > 0
    xmax = np.max(joints[:, :, 0], axis=-1, where=visible_joints, initial=joints[:, :, 0].min())
    xmin = np.min(joints[:, :, 0], axis=-1, where=visible_joints, initial=joints[:, :, 0].max())
    ymax = np.max(joints[:, :, 1], axis=-1, where=visible_joints, initial=joints[:, :, 1].min())
    ymin = np.min(joints[:, :, 1], axis=-1, where=visible_joints, initial=joints[:, :, 1].max())

    w = xmax - xmin
    h = ymax - ymin
    raw_area = w * h
    area = np.clip(raw_area, a_min=0, a_max=None) * (visible_joints.sum(axis=-1, keepdims=False) > 1)
    return area

filter_by_bbox_area(min_bbox_area)

Remove pose instances that has area of the corresponding bounding box less than a certain threshold.

Parameters:

Name Type Description Default
sample

Instance of PoseEstimationSample to modify. Modification done in-place.

 required
min_bbox_area Union[int, float]

Minimal bounding box area of the pose to keep.

 required

Returns:

Type Description
PoseEstimationSample

A pose sample after filtering.
Source code in V3_4/src/super_gradients/training/samples/pose_estimation_sample.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def filter_by_bbox_area(self, min_bbox_area: Union[int, float]) -> "PoseEstimationSample":
    """
    Remove pose instances that has area of the corresponding bounding box less than a certain threshold.

    :param sample:        Instance of PoseEstimationSample to modify. Modification done in-place.
    :param min_bbox_area: Minimal bounding box area of the pose to keep.
    :return:              A pose sample after filtering.
    """
    if self.bboxes_xywh is None:
        area = self.compute_area_of_joints_bounding_box(self.joints)
    else:
        area = self.bboxes_xywh[..., 2:4].prod(axis=-1)

    keep_mask = area >= min_bbox_area
    return self.filter_by_mask(keep_mask)

filter_by_mask(mask)

Remove pose instances with respect to given mask.

:remark: This is main method to modify instances of the sample. If you are implementing a subclass of PoseEstimationSample and adding extra field associated with each pose instance (Let's say you add a distance property for each pose from the camera), then you should override this method to do filtering on extra attribute as well.

Parameters:

Name Type Description Default
mask np.ndarray

A boolean or integer mask of samples to keep for given sample.

 required

Returns:

Type Description
PoseEstimationSample

A pose sample after filtering.
Source code in V3_4/src/super_gradients/training/samples/pose_estimation_sample.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def filter_by_mask(self, mask: np.ndarray) -> "PoseEstimationSample":
    """
    Remove pose instances with respect to given mask.

    :remark: This is main method to modify instances of the sample.
    If you are implementing a subclass of PoseEstimationSample and adding extra field associated with each pose
    instance (Let's say you add a distance property for each pose from the camera), then you should override
    this method to do filtering on extra attribute as well.

    :param mask:   A boolean or integer mask of samples to keep for given sample.
    :return:       A pose sample after filtering.
    """
    self.joints = self.joints[mask]
    self.is_crowd = self.is_crowd[mask]
    if self.bboxes_xywh is not None:
        self.bboxes_xywh = self.bboxes_xywh[mask]
    if self.areas is not None:
        self.areas = self.areas[mask]
    return self

filter_by_pose_area(min_instance_area)

Remove pose instances which area is less than a certain threshold.

Parameters:

Name Type Description Default
sample

Instance of PoseEstimationSample to modify. Modification done in-place.

 required
min_instance_area Union[int, float]

Minimal area of the pose to keep.

 required

Returns:

Type Description
PoseEstimationSample

A pose sample after filtering.
Source code in V3_4/src/super_gradients/training/samples/pose_estimation_sample.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
def filter_by_pose_area(self, min_instance_area: Union[int, float]) -> "PoseEstimationSample":
    """
    Remove pose instances which area is less than a certain threshold.

    :param sample:            Instance of PoseEstimationSample to modify. Modification done in-place.
    :param min_instance_area: Minimal area of the pose to keep.
    :return:                  A pose sample after filtering.
    """

    if self.areas is not None:
        areas = self.areas
    elif self.bboxes_xywh is not None:
        # 0.53 is a heuristic multiplier from COCO to approximate object area from bbox area
        areas = self.bboxes_xywh[..., 2:4].prod(axis=-1, keepdims=False) * 0.53
    else:
        areas = self.compute_area_of_joints_bounding_box(self.joints)

    keep_mask = areas >= min_instance_area
    return self.filter_by_mask(keep_mask)

filter_by_visible_joints(min_visible_joints)

Remove instances from the sample which has less than N visible joints.

Parameters:

Name Type Description Default
min_visible_joints int

A minimal number of visible joints a pose has to have in order to be kept.

 required

Returns:

Type Description
PoseEstimationSample

A pose sample after filtering.
Source code in V3_4/src/super_gradients/training/samples/pose_estimation_sample.py 
116
117
118
119
120
121
122
123
124
125
def filter_by_visible_joints(self, min_visible_joints: int) -> "PoseEstimationSample":
    """
    Remove instances from the sample which has less than N visible joints.

    :param min_visible_joints: A minimal number of visible joints a pose has to have in order to be kept.
    :return:                   A pose sample after filtering.
    """
    visible_joints_mask = self.joints[:, :, 2] > 0
    keep_mask: np.ndarray = np.sum(visible_joints_mask, axis=-1) >= min_visible_joints
    return self.filter_by_mask(keep_mask)

sanitize_sample()

Apply sanity checks on the pose sample, which includes: - Clamp bbox coordinates to ensure they are within image boundaries - Update visibility status of keypoints if they are outside of image boundaries - Update area if bbox clipping occurs This function does not remove instances, but may make them subject for removal instead.

Returns:

Type Description
PoseEstimationSample

self
Source code in V3_4/src/super_gradients/training/samples/pose_estimation_sample.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def sanitize_sample(self) -> "PoseEstimationSample":
    """
    Apply sanity checks on the pose sample, which includes:
    - Clamp bbox coordinates to ensure they are within image boundaries
    - Update visibility status of keypoints if they are outside of image boundaries
    - Update area if bbox clipping occurs
    This function does not remove instances, but may make them subject for removal instead.
    :return: self
    """
    image_height, image_width, _ = self.image.shape

    # Update joints visibility status
    outside_left = self.joints[:, :, 0] < 0
    outside_top = self.joints[:, :, 1] < 0
    outside_right = self.joints[:, :, 0] >= image_width
    outside_bottom = self.joints[:, :, 1] >= image_height

    outside_image_mask = outside_left | outside_top | outside_right | outside_bottom
    self.joints[outside_image_mask, 2] = 0

    if self.bboxes_xywh is not None:
        # Clamp bboxes to image boundaries
        clamped_boxes = xywh_to_xyxy(self.bboxes_xywh, image_shape=(image_height, image_width))
        clamped_boxes[..., [0, 2]] = np.clip(clamped_boxes[..., [0, 2]], 0, image_width - 1)
        clamped_boxes[..., [1, 3]] = np.clip(clamped_boxes[..., [1, 3]], 0, image_height - 1)
        clamped_boxes = xyxy_to_xywh(clamped_boxes, image_shape=(image_height, image_width))

        # Recompute sample areas if they are present
        if self.areas is not None:
            area_reduction_factor = clamped_boxes[..., 2:4].prod(axis=-1) / (self.bboxes_xywh[..., 2:4].prod(axis=-1) + 1e-6)
            self.areas = self.areas * area_reduction_factor

        self.bboxes_xywh = clamped_boxes
    return self