Skip to content

Register now for

# Reference for `ultralytics/models/yolo/pose/predict.py`

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/pose/predict.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!

## `` ultralytics.models.yolo.pose.predict.PosePredictor

    PosePredictor(cfg=DEFAULT_CFG, overrides=None, _callbacks=None)

Bases: `DetectionPredictor`

A class extending the DetectionPredictor class for prediction based on a pose model.

This class specializes in pose estimation, handling keypoints detection alongside standard object detection capabilities inherited from DetectionPredictor.

Attributes:

Name | Type | Description
--- | --- | ---
`args` | `namespace` | Configuration arguments for the predictor.
`model` | `Module` | The loaded YOLO pose model with keypoint detection capabilities.

Methods:

Name | Description
--- | ---
`construct_result` | Construct the result object from the prediction, including keypoints.

Examples:

    >>> from ultralytics.utils import ASSETS
    >>> from ultralytics.models.yolo.pose import PosePredictor
    >>> args = dict(model="yolo11n-pose.pt", source=ASSETS)
    >>> predictor = PosePredictor(overrides=args)
    >>> predictor.predict_cli()

Sets up a PosePredictor instance, configuring it for pose detection tasks and handling device-specific warnings for Apple MPS.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`cfg` | `Any` | Configuration for the predictor. | `DEFAULT_CFG`
`overrides` | `dict` | Configuration overrides that take precedence over cfg. | `None`
`_callbacks` | `list` | List of callback functions to be invoked during prediction. | `None`

Examples:

    >>> from ultralytics.utils import ASSETS
    >>> from ultralytics.models.yolo.pose import PosePredictor
    >>> args = dict(model="yolo11n-pose.pt", source=ASSETS)
    >>> predictor = PosePredictor(overrides=args)
    >>> predictor.predict_cli()

Source code in `ultralytics/models/yolo/pose/predict.py`

    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
    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
        """
        Initialize PosePredictor for pose estimation tasks.

        Sets up a PosePredictor instance, configuring it for pose detection tasks and handling device-specific
        warnings for Apple MPS.

        Args:
            cfg (Any): Configuration for the predictor.
            overrides (dict, optional): Configuration overrides that take precedence over cfg.
            _callbacks (list, optional): List of callback functions to be invoked during prediction.

        Examples:
            >>> from ultralytics.utils import ASSETS
            >>> from ultralytics.models.yolo.pose import PosePredictor
            >>> args = dict(model="yolo11n-pose.pt", source=ASSETS)
            >>> predictor = PosePredictor(overrides=args)
            >>> predictor.predict_cli()
        """
        super().__init__(cfg, overrides, _callbacks)
        self.args.task = "pose"
        if isinstance(self.args.device, str) and self.args.device.lower() == "mps":
            LOGGER.warning(
                "Apple MPS known Pose bug. Recommend 'device=cpu' for Pose models. "
                "See https://github.com/ultralytics/ultralytics/issues/4031."
            )

### `` construct_result

    construct_result(pred, img, orig_img, img_path)

Construct the result object from the prediction, including keypoints.

Extends the parent class implementation by extracting keypoint data from predictions and adding them to the result object.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`pred` | `Tensor` | The predicted bounding boxes, scores, and keypoints with shape (N, 6+K*D) where N is the number of detections, K is the number of keypoints, and D is the keypoint dimension. | required
`img` | `Tensor` | The processed input image tensor with shape (B, C, H, W). | required
`orig_img` | `ndarray` | The original unprocessed image as a numpy array. | required
`img_path` | `str` | The path to the original image file. | required

Returns:

Type | Description
--- | ---
`Results` | The result object containing the original image, image path, class names, bounding boxes, and keypoints.

Source code in `ultralytics/models/yolo/pose/predict.py`

    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
    def construct_result(self, pred, img, orig_img, img_path):
        """
        Construct the result object from the prediction, including keypoints.

        Extends the parent class implementation by extracting keypoint data from predictions and adding them to the
        result object.

        Args:
            pred (torch.Tensor): The predicted bounding boxes, scores, and keypoints with shape (N, 6+K*D) where N is
                the number of detections, K is the number of keypoints, and D is the keypoint dimension.
            img (torch.Tensor): The processed input image tensor with shape (B, C, H, W).
            orig_img (np.ndarray): The original unprocessed image as a numpy array.
            img_path (str): The path to the original image file.

        Returns:
            (Results): The result object containing the original image, image path, class names, bounding boxes, and
                keypoints.
        """
        result = super().construct_result(pred, img, orig_img, img_path)
        # Extract keypoints from prediction and reshape according to model's keypoint shape
        pred_kpts = pred[:, 6:].view(pred.shape[0], *self.model.kpt_shape)
        # Scale keypoints coordinates to match the original image dimensions
        pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
        result.update(keypoints=pred_kpts)
        return result

📅 Created 1 year ago ✏️ Updated 1 year ago