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
Bases: BasePredictor
RT-DETR (Real-Time Detection Transformer) Predictor extending the BasePredictor class for making predictions using Baidu's RT-DETR model.
This class leverages the power of 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.
Example
Attributes:
Name | Type | Description |
---|---|---|
imgsz | int | Image size for inference (must be square and scale-filled). |
args | dict | Argument overrides for the predictor. |
Parameters:
Name | Type | Description | Default |
---|---|---|---|
cfg | str | Path to a configuration file. Defaults to DEFAULT_CFG. | DEFAULT_CFG |
overrides | dict | Configuration overrides. Defaults to None. | None |
Source code in ultralytics/engine/predictor.py
postprocess
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
.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
preds | list | List of [predictions, extra] from the model. | required |
img | Tensor | Processed input images. | required |
orig_imgs | list or 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
pre_transform
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 scaleFilled.
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. |