Reference for ultralytics/models/rtdetr/predict.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/rtdetr/predict.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!
ultralytics.models.rtdetr.predict.RTDETRPredictor
RTDETRPredictor(cfg=DEFAULT_CFG, overrides=None, _callbacks=None)
Bases: BasePredictor
RT-DETR (Real-Time Detection Transformer) Predictor extending the BasePredictor class for making predictions.
This class leverages Vision Transformers to provide real-time object detection while maintaining high accuracy. It supports key features like efficient hybrid encoding and IoU-aware query selection.
Attributes:
Name | Type | Description |
---|---|---|
imgsz |
int
|
Image size for inference (must be square and scale-filled). |
args |
dict
|
Argument overrides for the predictor. |
model |
Module
|
The loaded RT-DETR model. |
batch |
list
|
Current batch of processed inputs. |
Examples:
>>> from ultralytics.utils import ASSETS
>>> from ultralytics.models.rtdetr import RTDETRPredictor
>>> args = dict(model="rtdetr-l.pt", source=ASSETS)
>>> predictor = RTDETRPredictor(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 |
|
postprocess
postprocess(preds, img, orig_imgs)
Postprocess the raw predictions from the model to generate bounding boxes and confidence scores.
The method filters detections based on confidence and class if specified in self.args
. It converts
model predictions to Results objects containing properly scaled bounding boxes.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
preds
|
List | Tuple
|
List of [predictions, extra] from the model, where predictions contain bounding boxes and scores. |
required |
img
|
Tensor
|
Processed input images with shape (N, 3, H, W). |
required |
orig_imgs
|
List | Tensor
|
Original, unprocessed images. |
required |
Returns:
Type | Description |
---|---|
List[Results]
|
A list of Results objects containing the post-processed bounding boxes, confidence scores, and class labels. |
Source code in ultralytics/models/rtdetr/predict.py
32 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 70 |
|
pre_transform
pre_transform(im)
Pre-transforms the input images before feeding them into the model for inference. The input images are letterboxed to ensure a square aspect ratio and scale-filled. The size must be square(640) and scale_filled.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
im
|
list[ndarray] | Tensor
|
Input images of shape (N,3,h,w) for tensor, [(h,w,3) x N] for list. |
required |
Returns:
Type | Description |
---|---|
list
|
List of pre-transformed images ready for model inference. |
Source code in ultralytics/models/rtdetr/predict.py
72 73 74 75 76 77 78 79 80 81 82 83 84 |
|