Skip to content

Reference for ultralytics/models/yolo/obb/predict.py

Note

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


ultralytics.models.yolo.obb.predict.OBBPredictor 

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

Bases: DetectionPredictor

A class extending the DetectionPredictor class for prediction based on an Oriented Bounding Box (OBB) model.

This predictor handles oriented bounding box detection tasks, processing images and returning results with rotated bounding boxes.

Attributes:

Name Type Description
args namespace

Configuration arguments for the predictor.
model Module

The loaded YOLO OBB model.

Examples:

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

This constructor sets up an OBBPredictor instance for oriented bounding box detection tasks.

Parameters:

Name Type Description Default
cfg dict

Default configuration for the predictor.

 DEFAULT_CFG
overrides dict

Configuration overrides that take precedence over the default config.

 None
_callbacks list

List of callback functions to be invoked during prediction.

 None

Examples:

>>> from ultralytics.utils import ASSETS
>>> from ultralytics.models.yolo.obb import OBBPredictor
>>> args = dict(model="yolo11n-obb.pt", source=ASSETS)
>>> predictor = OBBPredictor(overrides=args)
Source code in ultralytics/models/yolo/obb/predict.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initialize OBBPredictor with optional model and data configuration overrides.

    This constructor sets up an OBBPredictor instance for oriented bounding box detection tasks.

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

    Examples:
        >>> from ultralytics.utils import ASSETS
        >>> from ultralytics.models.yolo.obb import OBBPredictor
        >>> args = dict(model="yolo11n-obb.pt", source=ASSETS)
        >>> predictor = OBBPredictor(overrides=args)
    """
    super().__init__(cfg, overrides, _callbacks)
    self.args.task = "obb"

construct_result 

construct_result(pred, img, orig_img, img_path)

Construct the result object from the prediction.

Parameters:

Name Type Description Default
pred Tensor

The predicted bounding boxes, scores, and rotation angles with shape (N, 6) where the last dimension contains [x, y, w, h, confidence, class_id, angle].

 required
img Tensor

The image after preprocessing with shape (B, C, H, W).

 required
orig_img ndarray

The original image before preprocessing.

 required
img_path str

The path to the original image.

 required

Returns:

Type Description
Results

The result object containing the original image, image path, class names, and oriented bounding boxes.
Source code in ultralytics/models/yolo/obb/predict.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def construct_result(self, pred, img, orig_img, img_path):
    """
    Construct the result object from the prediction.

    Args:
        pred (torch.Tensor): The predicted bounding boxes, scores, and rotation angles with shape (N, 6) where
            the last dimension contains [x, y, w, h, confidence, class_id, angle].
        img (torch.Tensor): The image after preprocessing with shape (B, C, H, W).
        orig_img (np.ndarray): The original image before preprocessing.
        img_path (str): The path to the original image.

    Returns:
        (Results): The result object containing the original image, image path, class names, and oriented bounding boxes.
    """
    rboxes = ops.regularize_rboxes(torch.cat([pred[:, :4], pred[:, -1:]], dim=-1))
    rboxes[:, :4] = ops.scale_boxes(img.shape[2:], rboxes[:, :4], orig_img.shape, xywh=True)
    obb = torch.cat([rboxes, pred[:, 4:6]], dim=-1)
    return Results(orig_img, path=img_path, names=self.model.names, obb=obb)



📅 Created 1 year ago ✏️ Updated 7 months ago
jk4e glenn-jocher Burhan-Q