Skip to content

Samples

DepthEstimationSample dataclass

A dataclass representing a single depth estimation sample. Contains input image and depth map.

Parameters:

Name Type Description Default
image np.ndarray

np.ndarray, Image of [H, W, (C if colorful)] shape.

 required
depth_map np.ndarray

Depth map of [H, W] shape.

 required
Source code in V3_6/src/super_gradients/training/samples/depth_estimation_sample.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
@dataclasses.dataclass
class DepthEstimationSample:
    """
    A dataclass representing a single depth estimation sample.
    Contains input image and depth map.

    :param image:              np.ndarray, Image of [H, W, (C if colorful)] shape.
    :param depth_map:          Depth map of [H, W] shape.
    """

    __slots__ = ["image", "depth_map"]

    image: np.ndarray
    depth_map: np.ndarray

    def __init__(self, image: np.ndarray, depth_map: np.ndarray):
        # small sanity check
        dm_shape = depth_map.shape

        if len(dm_shape) == 3:
            if dm_shape[-1] == 1:
                depth_map = np.squeeze(depth_map, axis=-1)
            else:
                raise RuntimeError(f"Depth map should contain only H and W dimensions, got {len(dm_shape)} dimensions instead.")

        self.image = image
        self.depth_map = depth_map

DetectionSample dataclass

A data class describing a single object detection sample that comes from a dataset. It contains both input image and target information to train an object detection 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
bboxes_xywh

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

 required
labels np.ndarray

Numpy array of [N] shape with class label for each instance

 required
is_crowd Optional[np.ndarray]

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

 None
additional_samples Optional[List[DetectionSample]]

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

 None
Source code in V3_6/src/super_gradients/training/samples/detection_sample.py 
 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
@dataclasses.dataclass
class DetectionSample:
    """
    A data class describing a single object detection sample that comes from a dataset.
    It contains both input image and target information to train an object detection model.

    :param image:              Associated image with a sample. Can be in [H,W,C] or [C,H,W] format
    :param bboxes_xywh:        Numpy array of [N,4] shape with bounding box of each instance (XYWH)
    :param labels:             Numpy array of [N] shape with class label for each instance
    :param is_crowd:           (Optional) Numpy array of [N] shape with is_crowd flag for each instance
    :param additional_samples: (Optional) List of additional samples for the same image.
    """

    __slots__ = ["image", "bboxes_xyxy", "labels", "is_crowd", "additional_samples"]

    image: Union[np.ndarray, torch.Tensor]
    bboxes_xyxy: np.ndarray
    labels: np.ndarray
    is_crowd: np.ndarray
    additional_samples: Optional[List["DetectionSample"]]

    def __init__(
        self,
        image: Union[np.ndarray, torch.Tensor],
        bboxes_xyxy: np.ndarray,
        labels: np.ndarray,
        is_crowd: Optional[np.ndarray] = None,
        additional_samples: Optional[List["DetectionSample"]] = None,
    ):
        if is_crowd is None:
            is_crowd = np.zeros(len(labels), dtype=bool)

        if len(bboxes_xyxy) != len(labels):
            raise ValueError("Number of bounding boxes and labels must be equal. Got {len(bboxes_xyxy)} and {len(labels)} respectively")

        if len(bboxes_xyxy) != len(is_crowd):
            raise ValueError("Number of bounding boxes and is_crowd flags must be equal. Got {len(bboxes_xyxy)} and {len(is_crowd)} respectively")

        if len(bboxes_xyxy.shape) != 2 or bboxes_xyxy.shape[1] != 4:
            raise ValueError(f"Bounding boxes must be in [N,4] format. Shape of input bboxes is {bboxes_xyxy.shape}")

        if len(is_crowd.shape) != 1:
            raise ValueError(f"Number of is_crowd flags must be in [N] format. Shape of input is_crowd is {is_crowd.shape}")

        if len(labels.shape) != 1:
            raise ValueError("Labels must be in [N] format. Shape of input labels is {labels.shape}")

        self.image = image
        self.bboxes_xyxy = bboxes_xyxy
        self.labels = labels
        self.is_crowd = is_crowd
        self.additional_samples = additional_samples
        self.sanitize_sample()

    def sanitize_sample(self) -> "DetectionSample":
        """
        Apply sanity checks on the detection sample, which includes clamping of bounding boxes to image boundaries.
        This function does not remove instances, but may make them subject for removal later on.
        This method operates in-place and modifies the caller.
        :return: A DetectionSample after filtering (caller instance).
        """
        image_height, image_width = self.image.shape[:2]
        self.bboxes_xyxy = change_bbox_bounds_for_image_size(self.bboxes_xyxy, img_shape=(image_height, image_width))
        self.filter_by_bbox_area(0)
        return self

    def filter_by_mask(self, mask: np.ndarray) -> "DetectionSample":
        """
        Remove boxes & labels with respect to a given mask.
        This method operates in-place and modifies the caller.
        If you are implementing a subclass of DetectionSample and adding extra field associated with each bbox
        instance (Let's say you add a distance property for each bbox 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 DetectionSample after filtering (caller instance).
        """
        self.bboxes_xyxy = self.bboxes_xyxy[mask]
        self.labels = self.labels[mask]
        if self.is_crowd is not None:
            self.is_crowd = self.is_crowd[mask]
        return self

    def filter_by_bbox_area(self, min_bbox_area: Union[int, float]) -> "DetectionSample":
        """
        Remove pose instances that has area of the corresponding bounding box less than a certain threshold.
        This method operates in-place and modifies the caller.

        :param min_bbox_area: Minimal bounding box area of the pose to keep.
        :return:              A DetectionSample after filtering (caller instance).
        """
        bboxes_xywh = xyxy_to_xywh(self.bboxes_xyxy, image_shape=None)
        area = bboxes_xywh[..., 2:4].prod(axis=-1)
        keep_mask = area > min_bbox_area
        return self.filter_by_mask(keep_mask)

filter_by_bbox_area(min_bbox_area)

Remove pose instances that has area of the corresponding bounding box less than a certain threshold. This method operates in-place and modifies the caller.

Parameters:

Name Type Description Default
min_bbox_area Union[int, float]

Minimal bounding box area of the pose to keep.

 required

Returns:

Type Description
DetectionSample

A DetectionSample after filtering (caller instance).
Source code in V3_6/src/super_gradients/training/samples/detection_sample.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
def filter_by_bbox_area(self, min_bbox_area: Union[int, float]) -> "DetectionSample":
    """
    Remove pose instances that has area of the corresponding bounding box less than a certain threshold.
    This method operates in-place and modifies the caller.

    :param min_bbox_area: Minimal bounding box area of the pose to keep.
    :return:              A DetectionSample after filtering (caller instance).
    """
    bboxes_xywh = xyxy_to_xywh(self.bboxes_xyxy, image_shape=None)
    area = 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 boxes & labels with respect to a given mask. This method operates in-place and modifies the caller. If you are implementing a subclass of DetectionSample and adding extra field associated with each bbox instance (Let's say you add a distance property for each bbox 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
DetectionSample

A DetectionSample after filtering (caller instance).
Source code in V3_6/src/super_gradients/training/samples/detection_sample.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def filter_by_mask(self, mask: np.ndarray) -> "DetectionSample":
    """
    Remove boxes & labels with respect to a given mask.
    This method operates in-place and modifies the caller.
    If you are implementing a subclass of DetectionSample and adding extra field associated with each bbox
    instance (Let's say you add a distance property for each bbox 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 DetectionSample after filtering (caller instance).
    """
    self.bboxes_xyxy = self.bboxes_xyxy[mask]
    self.labels = self.labels[mask]
    if self.is_crowd is not None:
        self.is_crowd = self.is_crowd[mask]
    return self

sanitize_sample()

Apply sanity checks on the detection sample, which includes clamping of bounding boxes to image boundaries. This function does not remove instances, but may make them subject for removal later on. This method operates in-place and modifies the caller.

Returns:

Type Description
DetectionSample

A DetectionSample after filtering (caller instance).
Source code in V3_6/src/super_gradients/training/samples/detection_sample.py 
68
69
70
71
72
73
74
75
76
77
78
def sanitize_sample(self) -> "DetectionSample":
    """
    Apply sanity checks on the detection sample, which includes clamping of bounding boxes to image boundaries.
    This function does not remove instances, but may make them subject for removal later on.
    This method operates in-place and modifies the caller.
    :return: A DetectionSample after filtering (caller instance).
    """
    image_height, image_width = self.image.shape[:2]
    self.bboxes_xyxy = change_bbox_bounds_for_image_size(self.bboxes_xyxy, img_shape=(image_height, image_width))
    self.filter_by_bbox_area(0)
    return self

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_6/src/super_gradients/training/samples/pose_estimation_sample.py 
 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
162
@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 = change_bbox_bounds_for_image_size(clamped_boxes, img_shape=(image_height, image_width))
            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_6/src/super_gradients/training/samples/pose_estimation_sample.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
@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_6/src/super_gradients/training/samples/pose_estimation_sample.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
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_6/src/super_gradients/training/samples/pose_estimation_sample.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
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_6/src/super_gradients/training/samples/pose_estimation_sample.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
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_6/src/super_gradients/training/samples/pose_estimation_sample.py 
117
118
119
120
121
122
123
124
125
126
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_6/src/super_gradients/training/samples/pose_estimation_sample.py 
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
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 = change_bbox_bounds_for_image_size(clamped_boxes, img_shape=(image_height, image_width))
        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

SegmentationSample dataclass

A data class describing a single semantic segmentation sample that comes from a dataset. It contains both input image and target segmentation mask to train an semantic segmentation model.

Parameters:

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

Associated image with sample, in [H,W,C] (or H,W for greyscale) format.

 required
mask Union[np.ndarray, Image]

Associated segmentation mask with sample, in [H,W]

 required
Source code in V3_6/src/super_gradients/training/samples/segmentation_sample.py 
11
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
@dataclasses.dataclass
class SegmentationSample:
    """
    A data class describing a single semantic segmentation sample that comes from a dataset.
    It contains both input image and target segmentation mask to train an semantic segmentation model.

    :param image:              Associated image with sample, in [H,W,C] (or H,W for greyscale) format.
    :param mask:               Associated segmentation mask with sample, in [H,W]
    """

    __slots__ = ["image", "mask"]

    image: Union[np.ndarray, torch.Tensor]
    mask: Union[np.ndarray, torch.Tensor]

    def __init__(self, image: Union[np.ndarray, Image], mask: Union[np.ndarray, Image]):
        """
        Initialize segmentation sample
        :param image: Input image in [H,W,C] format. Can be numpy array or PIL image (in which case it will be converted to numpy array)
        :param mask:  Target mask in [H,W] format. Can be numpy array or PIL image (in which case it will be converted to numpy array)
        """
        if isinstance(image, Image):
            image = np.array(image)
        if isinstance(mask, Image):
            mask = np.array(mask)
        self.image = image
        self.mask = mask

__init__(image, mask)

Initialize segmentation sample

Parameters:

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

Input image in [H,W,C] format. Can be numpy array or PIL image (in which case it will be converted to numpy array)

 required
mask Union[np.ndarray, Image]

Target mask in [H,W] format. Can be numpy array or PIL image (in which case it will be converted to numpy array)

 required
Source code in V3_6/src/super_gradients/training/samples/segmentation_sample.py 
26
27
28
29
30
31
32
33
34
35
36
37
def __init__(self, image: Union[np.ndarray, Image], mask: Union[np.ndarray, Image]):
    """
    Initialize segmentation sample
    :param image: Input image in [H,W,C] format. Can be numpy array or PIL image (in which case it will be converted to numpy array)
    :param mask:  Target mask in [H,W] format. Can be numpy array or PIL image (in which case it will be converted to numpy array)
    """
    if isinstance(image, Image):
        image = np.array(image)
    if isinstance(mask, Image):
        mask = np.array(mask)
    self.image = image
    self.mask = mask