Skip to content

Register now for

# 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: dict[str, Any] | None = None,
        _callbacks: dict[str, list[callable]] | None = 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`

    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
    149
    150
    def __init__(
        self,
        cfg=DEFAULT_CFG,
        overrides: dict[str, Any] | None = None,
        _callbacks: dict[str, list[callable]] | None = 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 shortest 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 idx.shape[0] 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 = nms.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 1 year ago