Skip to content

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

Note

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


ultralytics.models.yolo.detect.predict.DetectionPredictor 

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

Bases: BasePredictor

A class extending the BasePredictor class for prediction based on a detection model.

This predictor specializes in object detection tasks, processing model outputs into meaningful detection results with bounding boxes and class predictions.

Attributes:

Name Type Description
args namespace

Configuration arguments for the predictor.
model Module

The detection model used for inference.
batch list

Batch of images and metadata for processing.

Methods:

Name Description
postprocess

Process raw model predictions into detection results.
construct_results

Build Results objects from processed predictions.
construct_result

Create a single Result object from a prediction.

Examples:

>>> from ultralytics.utils import ASSETS
>>> from ultralytics.models.yolo.detect import DetectionPredictor
>>> args = dict(model="yolo11n.pt", source=ASSETS)
>>> predictor = DetectionPredictor(overrides=args)
>>> predictor.predict_cli()
Source code in ultralytics/engine/predictor.py 
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
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initialize the BasePredictor class.

    Args:
        cfg (str | dict): Path to a configuration file or a configuration dictionary.
        overrides (dict | None): Configuration overrides.
        _callbacks (dict | None): Dictionary of callback functions.
    """
    self.args = get_cfg(cfg, overrides)
    self.save_dir = get_save_dir(self.args)
    if self.args.conf is None:
        self.args.conf = 0.25  # default conf=0.25
    self.done_warmup = False
    if self.args.show:
        self.args.show = check_imshow(warn=True)

    # Usable if setup is done
    self.model = None
    self.data = self.args.data  # data_dict
    self.imgsz = None
    self.device = None
    self.dataset = None
    self.vid_writer = {}  # dict of {save_path: video_writer, ...}
    self.plotted_img = None
    self.source_type = None
    self.seen = 0
    self.windows = []
    self.batch = None
    self.results = None
    self.transforms = None
    self.callbacks = _callbacks or callbacks.get_default_callbacks()
    self.txt_path = None
    self._lock = threading.Lock()  # for automatic thread-safe inference
    callbacks.add_integration_callbacks(self)

construct_result 

construct_result(pred, img, orig_img, img_path)

Construct a single Results object from one image prediction.

Parameters:

Name Type Description Default
pred Tensor

Predicted boxes and scores with shape (N, 6) where N is the number of detections.

 required
img Tensor

Preprocessed image tensor used for inference.

 required
orig_img ndarray

Original image before preprocessing.

 required
img_path str

Path to the original image file.

 required

Returns:

Type Description
Results

Results object containing the original image, image path, class names, and scaled bounding boxes.
Source code in ultralytics/models/yolo/detect/predict.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def construct_result(self, pred, img, orig_img, img_path):
    """
    Construct a single Results object from one image prediction.

    Args:
        pred (torch.Tensor): Predicted boxes and scores with shape (N, 6) where N is the number of detections.
        img (torch.Tensor): Preprocessed image tensor used for inference.
        orig_img (np.ndarray): Original image before preprocessing.
        img_path (str): Path to the original image file.

    Returns:
        (Results): Results object containing the original image, image path, class names, and scaled bounding boxes.
    """
    pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
    return Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6])

construct_results 

construct_results(preds, img, orig_imgs)

Construct a list of Results objects from model predictions.

Parameters:

Name Type Description Default
preds List[Tensor]

List of predicted bounding boxes and scores for each image.

 required
img Tensor

Batch of preprocessed images used for inference.

 required
orig_imgs List[ndarray]

List of original images before preprocessing.

 required

Returns:

Type Description
List[Results]

List of Results objects containing detection information for each image.
Source code in ultralytics/models/yolo/detect/predict.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def construct_results(self, preds, img, orig_imgs):
    """
    Construct a list of Results objects from model predictions.

    Args:
        preds (List[torch.Tensor]): List of predicted bounding boxes and scores for each image.
        img (torch.Tensor): Batch of preprocessed images used for inference.
        orig_imgs (List[np.ndarray]): List of original images before preprocessing.

    Returns:
        (List[Results]): List of Results objects containing detection information for each image.
    """
    return [
        self.construct_result(pred, img, orig_img, img_path)
        for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0])
    ]

postprocess 

postprocess(preds, img, orig_imgs, **kwargs)

Post-process predictions and return a list of Results objects.

This method applies non-maximum suppression to raw model predictions and prepares them for visualization and further analysis.

Parameters:

Name Type Description Default
preds Tensor

Raw predictions from the model.

 required
img Tensor

Processed input image tensor in model input format.

 required
orig_imgs Tensor | list

Original input images before preprocessing.

 required
**kwargs Any

Additional keyword arguments.

 {}

Returns:

Type Description
list

List of Results objects containing the post-processed predictions.

Examples:

>>> predictor = DetectionPredictor(overrides=dict(model="yolov8n.pt"))
>>> results = predictor.predict("path/to/image.jpg")
>>> processed_results = predictor.postprocess(preds, img, orig_imgs)
Source code in ultralytics/models/yolo/detect/predict.py 
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
def postprocess(self, preds, img, orig_imgs, **kwargs):
    """
    Post-process predictions and return a list of Results objects.

    This method applies non-maximum suppression to raw model predictions and prepares them for visualization and
    further analysis.

    Args:
        preds (torch.Tensor): Raw predictions from the model.
        img (torch.Tensor): Processed input image tensor in model input format.
        orig_imgs (torch.Tensor | list): Original input images before preprocessing.
        **kwargs (Any): Additional keyword arguments.

    Returns:
        (list): List of Results objects containing the post-processed predictions.

    Examples:
        >>> predictor = DetectionPredictor(overrides=dict(model="yolov8n.pt"))
        >>> results = predictor.predict("path/to/image.jpg")
        >>> processed_results = predictor.postprocess(preds, img, orig_imgs)
    """
    preds = ops.non_max_suppression(
        preds,
        self.args.conf,
        self.args.iou,
        self.args.classes,
        self.args.agnostic_nms,
        max_det=self.args.max_det,
        nc=len(self.model.names),
        end2end=getattr(self.model, "end2end", False),
        rotated=self.args.task == "obb",
    )

    if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
        orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

    return self.construct_results(preds, img, orig_imgs, **kwargs)



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