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: Optional[Dict[str, Any]] = None,
    _callbacks: Optional[Dict[str, List[callable]]] = 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.

get_obj_feats

Extract object features from the feature maps.

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
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
def __init__(
    self,
    cfg=DEFAULT_CFG,
    overrides: Optional[Dict[str, Any]] = None,
    _callbacks: Optional[Dict[str, List[callable]]] = None,
):
    """
    Initialize the BasePredictor class.

    Args:
        cfg (str | dict): Path to a configuration file or a configuration dictionary.
        overrides (dict, optional): Configuration overrides.
        _callbacks (dict, optional): 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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
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
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
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])
    ]

get_obj_feats

get_obj_feats(feat_maps, idxs)

Extract object features from the feature maps.

Source code in ultralytics/models/yolo/detect/predict.py
84
85
86
87
88
89
90
91
92
def get_obj_feats(self, feat_maps, idxs):
    """Extract object features from the feature maps."""
    import torch

    s = min([x.shape[1] for x in feat_maps])  # find smallest vector length
    obj_feats = torch.cat(
        [x.permute(0, 2, 3, 1).reshape(x.shape[0], -1, s, x.shape[1] // s).mean(dim=-1) for x in feat_maps], dim=1
    )  # mean reduce all vectors to same length
    return [feats[idx] if len(idx) else [] for feats, idx in zip(obj_feats, idxs)]  # for each img in batch

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="yolo11n.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
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
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="yolo11n.pt"))
        >>> results = predictor.predict("path/to/image.jpg")
        >>> processed_results = predictor.postprocess(preds, img, orig_imgs)
    """
    save_feats = getattr(self, "_feats", None) is not None
    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=0 if self.args.task == "detect" else len(self.model.names),
        end2end=getattr(self.model, "end2end", False),
        rotated=self.args.task == "obb",
        return_idxs=save_feats,
    )

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

    if save_feats:
        obj_feats = self.get_obj_feats(self._feats, preds[1])
        preds = preds[0]

    results = self.construct_results(preds, img, orig_imgs, **kwargs)

    if save_feats:
        for r, f in zip(results, obj_feats):
            r.feats = f  # add object features to results

    return results





📅 Created 1 year ago ✏️ Updated 8 months ago