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()
Source code in ultralytics/models/yolo/obb/predict.py 
29
30
31
32
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """Initialize OBBPredictor with optional model and data configuration overrides."""
    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 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
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 6 months ago
jk4e glenn-jocher Burhan-Q