dgenerate-ultralytics-headless 8.3.196__py3-none-any.whl → 8.3.248__py3-none-any.whl

ultralytics/hub/session.py

@@ -19,8 +19,7 @@ AGENT_NAME = f"python-{__version__}-colab" if IS_COLAB else f"python-{__version_
 
 
 class HUBTrainingSession:
-    """
-    HUB training session for Ultralytics HUB YOLO models.
+    """HUB training session for Ultralytics HUB YOLO models.
 
     This class encapsulates the functionality for interacting with Ultralytics HUB during model training, including
     model creation, metrics tracking, and checkpoint uploading.
@@ -45,12 +44,11 @@ class HUBTrainingSession:
     """
 
     def __init__(self, identifier: str):
-        """
-        Initialize the HUBTrainingSession with the provided model identifier.
+        """Initialize the HUBTrainingSession with the provided model identifier.
 
         Args:
-            identifier (str): Model identifier used to initialize the HUB training session. It can be a URL string
-                or a model key with specific format.
+            identifier (str): Model identifier used to initialize the HUB training session. It can be a URL string or a
+                model key with specific format.
 
         Raises:
             ValueError: If the provided model identifier is invalid.
@@ -93,8 +91,7 @@ class HUBTrainingSession:
 
     @classmethod
     def create_session(cls, identifier: str, args: dict[str, Any] | None = None):
-        """
-        Create an authenticated HUBTrainingSession or return None.
+        """Create an authenticated HUBTrainingSession or return None.
 
         Args:
             identifier (str): Model identifier used to initialize the HUB training session.
@@ -114,8 +111,7 @@ class HUBTrainingSession:
             return None
 
     def load_model(self, model_id: str):
-        """
-        Load an existing model from Ultralytics HUB using the provided model identifier.
+        """Load an existing model from Ultralytics HUB using the provided model identifier.
 
         Args:
             model_id (str): The identifier of the model to load.
@@ -140,8 +136,7 @@ class HUBTrainingSession:
         LOGGER.info(f"{PREFIX}View model at {self.model_url} 🚀")
 
     def create_model(self, model_args: dict[str, Any]):
-        """
-        Initialize a HUB training session with the specified model arguments.
+        """Initialize a HUB training session with the specified model arguments.
 
         Args:
             model_args (dict[str, Any]): Arguments for creating the model, including batch size, epochs, image size,
@@ -186,8 +181,7 @@ class HUBTrainingSession:
 
     @staticmethod
     def _parse_identifier(identifier: str):
-        """
-        Parse the given identifier to determine the type and extract relevant components.
+        """Parse the given identifier to determine the type and extract relevant components.
 
         The method supports different identifier formats:
             - A HUB model URL https://hub.ultralytics.com/models/MODEL
@@ -218,12 +212,11 @@ class HUBTrainingSession:
         return api_key, model_id, filename
 
     def _set_train_args(self):
-        """
-        Initialize training arguments and create a model entry on the Ultralytics HUB.
+        """Initialize training arguments and create a model entry on the Ultralytics HUB.
 
-        This method sets up training arguments based on the model's state and updates them with any additional
-        arguments provided. It handles different states of the model, such as whether it's resumable, pretrained,
-        or requires specific file setup.
+        This method sets up training arguments based on the model's state and updates them with any additional arguments
+        provided. It handles different states of the model, such as whether it's resumable, pretrained, or requires
+        specific file setup.
 
         Raises:
             ValueError: If the model is already trained, if required dataset information is missing, or if there are
@@ -261,8 +254,7 @@ class HUBTrainingSession:
         *args,
         **kwargs,
     ):
-        """
-        Execute request_func with retries, timeout handling, optional threading, and progress tracking.
+        """Execute request_func with retries, timeout handling, optional threading, and progress tracking.
 
         Args:
             request_func (callable): The function to execute.
@@ -342,8 +334,7 @@ class HUBTrainingSession:
         return status_code in retry_codes
 
     def _get_failure_message(self, response, retry: int, timeout: int) -> str:
-        """
-        Generate a retry message based on the response status code.
+        """Generate a retry message based on the response status code.
 
         Args:
             response (requests.Response): The HTTP response object.
@@ -379,8 +370,7 @@ class HUBTrainingSession:
         map: float = 0.0,
         final: bool = False,
     ) -> None:
-        """
-        Upload a model checkpoint to Ultralytics HUB.
+        """Upload a model checkpoint to Ultralytics HUB.
 
         Args:
             epoch (int): The current training epoch.

ultralytics/hub/utils.py

@@ -21,8 +21,7 @@ HELP_MSG = "If this issue persists please visit https://github.com/ultralytics/h
 
 
 def request_with_credentials(url: str) -> Any:
-    """
-    Make an AJAX request with cookies attached in a Google Colab environment.
+    """Make an AJAX request with cookies attached in a Google Colab environment.
 
     Args:
         url (str): The URL to make the request to.
@@ -35,8 +34,8 @@ def request_with_credentials(url: str) -> Any:
     """
     if not IS_COLAB:
         raise OSError("request_with_credentials() must run in a Colab environment")
-    from google.colab import output  # noqa
-    from IPython import display  # noqa
+    from google.colab import output
+    from IPython import display
 
     display.display(
         display.Javascript(
@@ -62,8 +61,7 @@ def request_with_credentials(url: str) -> Any:
 
 
 def requests_with_progress(method: str, url: str, **kwargs):
-    """
-    Make an HTTP request using the specified method and URL, with an optional progress bar.
+    """Make an HTTP request using the specified method and URL, with an optional progress bar.
 
     Args:
         method (str): The HTTP method to use (e.g. 'GET', 'POST').
@@ -106,8 +104,7 @@ def smart_request(
     progress: bool = False,
     **kwargs,
 ):
-    """
-    Make an HTTP request using the 'requests' library, with exponential backoff retries up to a specified timeout.
+    """Make an HTTP request using the 'requests' library, with exponential backoff retries up to a specified timeout.
 
     Args:
         method (str): The HTTP method to use for the request. Choices are 'post' and 'get'.

ultralytics/models/__init__.py

@@ -6,4 +6,4 @@ from .rtdetr import RTDETR
 from .sam import SAM
 from .yolo import YOLO, YOLOE, YOLOWorld
 
-__all__ = "YOLO", "RTDETR", "SAM", "FastSAM", "NAS", "YOLOWorld", "YOLOE"  # allow simpler import
+__all__ = "NAS", "RTDETR", "SAM", "YOLO", "YOLOE", "FastSAM", "YOLOWorld"  # allow simpler import

ultralytics/models/fastsam/__init__.py

@@ -4,4 +4,4 @@ from .model import FastSAM
 from .predict import FastSAMPredictor
 from .val import FastSAMValidator
 
-__all__ = "FastSAMPredictor", "FastSAM", "FastSAMValidator"
+__all__ = "FastSAM", "FastSAMPredictor", "FastSAMValidator"

ultralytics/models/fastsam/model.py

@@ -12,8 +12,7 @@ from .val import FastSAMValidator
 
 
 class FastSAM(Model):
-    """
-    FastSAM model interface for segment anything tasks.
+    """FastSAM model interface for Segment Anything tasks.
 
     This class extends the base Model class to provide specific functionality for the FastSAM (Fast Segment Anything
     Model) implementation, allowing for efficient and accurate image segmentation with optional prompting support.
@@ -36,11 +35,11 @@ class FastSAM(Model):
         >>> results = model.predict("image.jpg", bboxes=[[100, 100, 200, 200]])
     """
 
-    def __init__(self, model: str = "FastSAM-x.pt"):
+    def __init__(self, model: str | Path = "FastSAM-x.pt"):
         """Initialize the FastSAM model with the specified pre-trained weights."""
         if str(model) == "FastSAM.pt":
             model = "FastSAM-x.pt"
-        assert Path(model).suffix not in {".yaml", ".yml"}, "FastSAM models only support pre-trained models."
+        assert Path(model).suffix not in {".yaml", ".yml"}, "FastSAM only supports pre-trained weights."
         super().__init__(model=model, task="segment")
 
     def predict(
@@ -53,15 +52,14 @@ class FastSAM(Model):
         texts: list | None = None,
         **kwargs: Any,
     ):
-        """
-        Perform segmentation prediction on image or video source.
+        """Perform segmentation prediction on image or video source.
 
-        Supports prompted segmentation with bounding boxes, points, labels, and texts. The method packages these
-        prompts and passes them to the parent class predict method for processing.
+        Supports prompted segmentation with bounding boxes, points, labels, and texts. The method packages these prompts
+        and passes them to the parent class predict method for processing.
 
         Args:
-            source (str | PIL.Image | np.ndarray): Input source for prediction, can be a file path, URL, PIL image,
-                or numpy array.
+            source (str | PIL.Image | np.ndarray): Input source for prediction, can be a file path, URL, PIL image, or
+                numpy array.
             stream (bool): Whether to enable real-time streaming mode for video inputs.
             bboxes (list, optional): Bounding box coordinates for prompted segmentation in format [[x1, y1, x2, y2]].
             points (list, optional): Point coordinates for prompted segmentation in format [[x, y]].

ultralytics/models/fastsam/predict.py

@@ -4,16 +4,16 @@ import torch
 from PIL import Image
 
 from ultralytics.models.yolo.segment import SegmentationPredictor
-from ultralytics.utils import DEFAULT_CFG, checks
+from ultralytics.utils import DEFAULT_CFG
 from ultralytics.utils.metrics import box_iou
 from ultralytics.utils.ops import scale_masks
+from ultralytics.utils.torch_utils import TORCH_1_10
 
 from .utils import adjust_bboxes_to_image_border
 
 
 class FastSAMPredictor(SegmentationPredictor):
-    """
-    FastSAMPredictor is specialized for fast SAM (Segment Anything Model) segmentation prediction tasks.
+    """FastSAMPredictor is specialized for fast SAM (Segment Anything Model) segmentation prediction tasks.
 
     This class extends the SegmentationPredictor, customizing the prediction pipeline specifically for fast SAM. It
     adjusts post-processing steps to incorporate mask prediction and non-maximum suppression while optimizing for
@@ -22,8 +22,7 @@ class FastSAMPredictor(SegmentationPredictor):
     Attributes:
         prompts (dict): Dictionary containing prompt information for segmentation (bboxes, points, labels, texts).
         device (torch.device): Device on which model and tensors are processed.
-        clip_model (Any, optional): CLIP model for text-based prompting, loaded on demand.
-        clip_preprocess (Any, optional): CLIP preprocessing function for images, loaded on demand.
+        clip (Any, optional): CLIP model used for text-based prompting, loaded on demand.
 
     Methods:
         postprocess: Apply postprocessing to FastSAM predictions and handle prompts.
@@ -32,8 +31,7 @@ class FastSAMPredictor(SegmentationPredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """
-        Initialize the FastSAMPredictor with configuration and callbacks.
+        """Initialize the FastSAMPredictor with configuration and callbacks.
 
         This initializes a predictor specialized for Fast SAM (Segment Anything Model) segmentation tasks. The predictor
         extends SegmentationPredictor with custom post-processing for mask prediction and non-maximum suppression
@@ -48,8 +46,7 @@ class FastSAMPredictor(SegmentationPredictor):
         self.prompts = {}
 
     def postprocess(self, preds, img, orig_imgs):
-        """
-        Apply postprocessing to FastSAM predictions and handle prompts.
+        """Apply postprocessing to FastSAM predictions and handle prompts.
 
         Args:
             preds (list[torch.Tensor]): Raw predictions from the model.
@@ -76,8 +73,7 @@ class FastSAMPredictor(SegmentationPredictor):
         return self.prompt(results, bboxes=bboxes, points=points, labels=labels, texts=texts)
 
     def prompt(self, results, bboxes=None, points=None, labels=None, texts=None):
-        """
-        Perform image segmentation inference based on cues like bounding boxes, points, and text prompts.
+        """Perform image segmentation inference based on cues like bounding boxes, points, and text prompts.
 
         Args:
             results (Results | list[Results]): Original inference results from FastSAM models without any prompts.
@@ -100,7 +96,7 @@ class FastSAMPredictor(SegmentationPredictor):
                 continue
             masks = result.masks.data
             if masks.shape[1:] != result.orig_shape:
-                masks = scale_masks(masks[None], result.orig_shape)[0]
+                masks = (scale_masks(masks[None].float(), result.orig_shape)[0] > 0.5).byte()
             # bboxes prompt
             idx = torch.zeros(len(result), dtype=torch.bool, device=self.device)
             if bboxes is not None:
@@ -119,7 +115,7 @@ class FastSAMPredictor(SegmentationPredictor):
                     labels = torch.ones(points.shape[0])
                 labels = torch.as_tensor(labels, dtype=torch.int32, device=self.device)
                 assert len(labels) == len(points), (
-                    f"Expected `labels` with same size as `point`, but got {len(labels)} and {len(points)}"
+                    f"Expected `labels` to have the same length as `points`, but got {len(labels)} and {len(points)}."
                 )
                 point_idx = (
                     torch.ones(len(result), dtype=torch.bool, device=self.device)
@@ -135,7 +131,7 @@ class FastSAMPredictor(SegmentationPredictor):
                 crop_ims, filter_idx = [], []
                 for i, b in enumerate(result.boxes.xyxy.tolist()):
                     x1, y1, x2, y2 = (int(x) for x in b)
-                    if masks[i].sum() <= 100:
+                    if (masks[i].sum() if TORCH_1_10 else masks[i].sum(0).sum()) <= 100:  # torch 1.9 bug workaround
                         filter_idx.append(i)
                         continue
                     crop_ims.append(Image.fromarray(result.orig_img[y1:y2, x1:x2, ::-1]))
@@ -150,8 +146,7 @@ class FastSAMPredictor(SegmentationPredictor):
         return prompt_results
 
     def _clip_inference(self, images, texts):
-        """
-        Perform CLIP inference to calculate similarity between images and text prompts.
+        """Perform CLIP inference to calculate similarity between images and text prompts.
 
         Args:
             images (list[PIL.Image]): List of source images, each should be PIL.Image with RGB channel order.
@@ -160,20 +155,14 @@ class FastSAMPredictor(SegmentationPredictor):
         Returns:
             (torch.Tensor): Similarity matrix between given images and texts with shape (M, N).
         """
-        try:
-            import clip
-        except ImportError:
-            checks.check_requirements("git+https://github.com/ultralytics/CLIP.git")
-            import clip
-        if (not hasattr(self, "clip_model")) or (not hasattr(self, "clip_preprocess")):
-            self.clip_model, self.clip_preprocess = clip.load("ViT-B/32", device=self.device)
-        images = torch.stack([self.clip_preprocess(image).to(self.device) for image in images])
-        tokenized_text = clip.tokenize(texts).to(self.device)
-        image_features = self.clip_model.encode_image(images)
-        text_features = self.clip_model.encode_text(tokenized_text)
-        image_features /= image_features.norm(dim=-1, keepdim=True)  # (N, 512)
-        text_features /= text_features.norm(dim=-1, keepdim=True)  # (M, 512)
-        return (image_features * text_features[:, None]).sum(-1)  # (M, N)
+        from ultralytics.nn.text_model import CLIP
+
+        if not hasattr(self, "clip"):
+            self.clip = CLIP("ViT-B/32", device=self.device)
+        images = torch.stack([self.clip.image_preprocess(image).to(self.device) for image in images])
+        image_features = self.clip.encode_image(images)
+        text_features = self.clip.encode_text(self.clip.tokenize(texts))
+        return text_features @ image_features.T  # (M, N)
 
     def set_prompts(self, prompts):
         """Set prompts to be used during inference."""

ultralytics/models/fastsam/utils.py

@@ -2,8 +2,7 @@
 
 
 def adjust_bboxes_to_image_border(boxes, image_shape, threshold=20):
-    """
-    Adjust bounding boxes to stick to image border if they are within a certain threshold.
+    """Adjust bounding boxes to stick to image border if they are within a certain threshold.
 
     Args:
         boxes (torch.Tensor): Bounding boxes with shape (N, 4) in xyxy format.

ultralytics/models/fastsam/val.py

@@ -4,10 +4,9 @@ from ultralytics.models.yolo.segment import SegmentationValidator
 
 
 class FastSAMValidator(SegmentationValidator):
-    """
-    Custom validation class for Fast SAM (Segment Anything Model) segmentation in Ultralytics YOLO framework.
+    """Custom validation class for FastSAM (Segment Anything Model) segmentation in the Ultralytics YOLO framework.
 
-    Extends the SegmentationValidator class, customizing the validation process specifically for Fast SAM. This class
+    Extends the SegmentationValidator class, customizing the validation process specifically for FastSAM. This class
     sets the task to 'segment' and uses the SegmentMetrics for evaluation. Additionally, plotting features are disabled
     to avoid errors during validation.
 
@@ -19,15 +18,14 @@ class FastSAMValidator(SegmentationValidator):
         metrics (SegmentMetrics): Segmentation metrics calculator for evaluation.
 
     Methods:
-        __init__: Initialize the FastSAMValidator with custom settings for Fast SAM.
+        __init__: Initialize the FastSAMValidator with custom settings for FastSAM.
     """
 
     def __init__(self, dataloader=None, save_dir=None, args=None, _callbacks=None):
-        """
-        Initialize the FastSAMValidator class, setting the task to 'segment' and metrics to SegmentMetrics.
+        """Initialize the FastSAMValidator class, setting the task to 'segment' and metrics to SegmentMetrics.
 
         Args:
-            dataloader (torch.utils.data.DataLoader, optional): Dataloader to be used for validation.
+            dataloader (torch.utils.data.DataLoader, optional): DataLoader to be used for validation.
             save_dir (Path, optional): Directory to save results.
             args (SimpleNamespace, optional): Configuration for the validator.
             _callbacks (list, optional): List of callback functions to be invoked during validation.

ultralytics/models/nas/__init__.py

@@ -4,4 +4,4 @@ from .model import NAS
 from .predict import NASPredictor
 from .val import NASValidator
 
-__all__ = "NASPredictor", "NASValidator", "NAS"
+__all__ = "NAS", "NASPredictor", "NASValidator"

ultralytics/models/nas/model.py

@@ -18,11 +18,10 @@ from .val import NASValidator
 
 
 class NAS(Model):
-    """
-    YOLO-NAS model for object detection.
+    """YOLO-NAS model for object detection.
 
-    This class provides an interface for the YOLO-NAS models and extends the `Model` class from Ultralytics engine.
-    It is designed to facilitate the task of object detection using pre-trained or custom-trained YOLO-NAS models.
+    This class provides an interface for the YOLO-NAS models and extends the `Model` class from Ultralytics engine. It
+    is designed to facilitate the task of object detection using pre-trained or custom-trained YOLO-NAS models.
 
     Attributes:
         model (torch.nn.Module): The loaded YOLO-NAS model.
@@ -48,8 +47,7 @@ class NAS(Model):
         super().__init__(model, task="detect")
 
     def _load(self, weights: str, task=None) -> None:
-        """
-        Load an existing NAS model weights or create a new NAS model with pretrained weights.
+        """Load an existing NAS model weights or create a new NAS model with pretrained weights.
 
         Args:
             weights (str): Path to the model weights file or model name.
@@ -83,8 +81,7 @@ class NAS(Model):
         self.model.eval()
 
     def info(self, detailed: bool = False, verbose: bool = True) -> dict[str, Any]:
-        """
-        Log model information.
+        """Log model information.
 
         Args:
             detailed (bool): Show detailed information about model.

ultralytics/models/nas/predict.py

@@ -7,12 +7,11 @@ from ultralytics.utils import ops
 
 
 class NASPredictor(DetectionPredictor):
-    """
-    Ultralytics YOLO NAS Predictor for object detection.
+    """Ultralytics YOLO NAS Predictor for object detection.
 
-    This class extends the DetectionPredictor from Ultralytics engine and is responsible for post-processing the
-    raw predictions generated by the YOLO NAS models. It applies operations like non-maximum suppression and
-    scaling the bounding boxes to fit the original image dimensions.
+    This class extends the DetectionPredictor from Ultralytics engine and is responsible for post-processing the raw
+    predictions generated by the YOLO NAS models. It applies operations like non-maximum suppression and scaling the
+    bounding boxes to fit the original image dimensions.
 
     Attributes:
         args (Namespace): Namespace containing various configurations for post-processing including confidence
@@ -33,12 +32,11 @@ class NASPredictor(DetectionPredictor):
     """
 
     def postprocess(self, preds_in, img, orig_imgs):
-        """
-        Postprocess NAS model predictions to generate final detection results.
+        """Postprocess NAS model predictions to generate final detection results.
 
         This method takes raw predictions from a YOLO NAS model, converts bounding box formats, and applies
-        post-processing operations to generate the final detection results compatible with Ultralytics
-        result visualization and analysis tools.
+        post-processing operations to generate the final detection results compatible with Ultralytics result
+        visualization and analysis tools.
 
         Args:
             preds_in (list): Raw predictions from the NAS model, typically containing bounding boxes and class scores.

ultralytics/models/nas/val.py

@@ -9,8 +9,7 @@ __all__ = ["NASValidator"]
 
 
 class NASValidator(DetectionValidator):
-    """
-    Ultralytics YOLO NAS Validator for object detection.
+    """Ultralytics YOLO NAS Validator for object detection.
 
     Extends DetectionValidator from the Ultralytics models package and is designed to post-process the raw predictions
     generated by YOLO NAS models. It performs non-maximum suppression to remove overlapping and low-confidence boxes,

ultralytics/models/rtdetr/__init__.py

@@ -4,4 +4,4 @@ from .model import RTDETR
 from .predict import RTDETRPredictor
 from .val import RTDETRValidator
 
-__all__ = "RTDETRPredictor", "RTDETRValidator", "RTDETR"
+__all__ = "RTDETR", "RTDETRPredictor", "RTDETRValidator"

ultralytics/models/rtdetr/model.py

@@ -11,6 +11,7 @@ References:
 
 from ultralytics.engine.model import Model
 from ultralytics.nn.tasks import RTDETRDetectionModel
+from ultralytics.utils.torch_utils import TORCH_1_11
 
 from .predict import RTDETRPredictor
 from .train import RTDETRTrainer
@@ -18,11 +19,10 @@ from .val import RTDETRValidator
 
 
 class RTDETR(Model):
-    """
-    Interface for Baidu's RT-DETR model, a Vision Transformer-based real-time object detector.
+    """Interface for Baidu's RT-DETR model, a Vision Transformer-based real-time object detector.
 
-    This model provides real-time performance with high accuracy. It supports efficient hybrid encoding, IoU-aware
-    query selection, and adaptable inference speed.
+    This model provides real-time performance with high accuracy. It supports efficient hybrid encoding, IoU-aware query
+    selection, and adaptable inference speed.
 
     Attributes:
         model (str): Path to the pre-trained model.
@@ -38,18 +38,17 @@ class RTDETR(Model):
     """
 
     def __init__(self, model: str = "rtdetr-l.pt") -> None:
-        """
-        Initialize the RT-DETR model with the given pre-trained model file.
+        """Initialize the RT-DETR model with the given pre-trained model file.
 
         Args:
             model (str): Path to the pre-trained model. Supports .pt, .yaml, and .yml formats.
         """
+        assert TORCH_1_11, "RTDETR requires torch>=1.11"
         super().__init__(model=model, task="detect")
 
     @property
     def task_map(self) -> dict:
-        """
-        Return a task map for RT-DETR, associating tasks with corresponding Ultralytics classes.
+        """Return a task map for RT-DETR, associating tasks with corresponding Ultralytics classes.
 
         Returns:
             (dict): A dictionary mapping task names to Ultralytics task classes for the RT-DETR model.

ultralytics/models/rtdetr/predict.py

@@ -9,11 +9,10 @@ from ultralytics.utils import ops
 
 
 class RTDETRPredictor(BasePredictor):
-    """
-    RT-DETR (Real-Time Detection Transformer) Predictor extending the BasePredictor class for making predictions.
+    """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.
+    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:
         imgsz (int): Image size for inference (must be square and scale-filled).
@@ -34,21 +33,20 @@ class RTDETRPredictor(BasePredictor):
     """
 
     def postprocess(self, preds, img, orig_imgs):
-        """
-        Postprocess the raw predictions from the model to generate bounding boxes and confidence scores.
+        """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.
+        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.
 
         Args:
-            preds (list | tuple): List of [predictions, extra] from the model, where predictions contain
-                bounding boxes and scores.
+            preds (list | tuple): List of [predictions, extra] from the model, where predictions contain bounding boxes
+                and scores.
             img (torch.Tensor): Processed input images with shape (N, 3, H, W).
             orig_imgs (list | torch.Tensor): Original, unprocessed images.
 
         Returns:
-            results (list[Results]): A list of Results objects containing the post-processed bounding boxes,
-                confidence scores, and class labels.
+            results (list[Results]): A list of Results objects containing the post-processed bounding boxes, confidence
+                scores, and class labels.
         """
         if not isinstance(preds, (list, tuple)):  # list for PyTorch inference but list[0] Tensor for export inference
             preds = [preds, None]
@@ -57,7 +55,7 @@ class RTDETRPredictor(BasePredictor):
         bboxes, scores = preds[0].split((4, nd - 4), dim=-1)
 
         if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
-            orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
+            orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)[..., ::-1]
 
         results = []
         for bbox, score, orig_img, img_path in zip(bboxes, scores, orig_imgs, self.batch[0]):  # (300, 4)
@@ -75,15 +73,13 @@ class RTDETRPredictor(BasePredictor):
         return results
 
     def pre_transform(self, im):
-        """
-        Pre-transform input images before feeding them into the model for inference.
+        """Pre-transform 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.
+        The input images are letterboxed to ensure a square aspect ratio and scale-filled.
 
         Args:
-            im (list[np.ndarray]  | torch.Tensor): Input images of shape (N, 3, H, W) for tensor,
-                [(H, W, 3) x N] for list.
+            im (list[np.ndarray] | torch.Tensor): Input images of shape (N, 3, H, W) for tensor, [(H, W, 3) x N] for
+                list.
 
         Returns:
             (list): List of pre-transformed images ready for model inference.

ultralytics/models/rtdetr/train.py

@@ -12,12 +12,11 @@ from .val import RTDETRDataset, RTDETRValidator
 
 
 class RTDETRTrainer(DetectionTrainer):
-    """
-    Trainer class for the RT-DETR model developed by Baidu for real-time object detection.
+    """Trainer class for the RT-DETR model developed by Baidu for real-time object detection.
 
-    This class extends the DetectionTrainer class for YOLO to adapt to the specific features and architecture of RT-DETR.
-    The model leverages Vision Transformers and has capabilities like IoU-aware query selection and adaptable inference
-    speed.
+    This class extends the DetectionTrainer class for YOLO to adapt to the specific features and architecture of
+    RT-DETR. The model leverages Vision Transformers and has capabilities like IoU-aware query selection and adaptable
+    inference speed.
 
     Attributes:
         loss_names (tuple): Names of the loss components used for training.
@@ -31,20 +30,19 @@ class RTDETRTrainer(DetectionTrainer):
         build_dataset: Build and return an RT-DETR dataset for training or validation.
         get_validator: Return a DetectionValidator suitable for RT-DETR model validation.
 
-    Notes:
-        - F.grid_sample used in RT-DETR does not support the `deterministic=True` argument.
-        - AMP training can lead to NaN outputs and may produce errors during bipartite graph matching.
-
     Examples:
         >>> from ultralytics.models.rtdetr.train import RTDETRTrainer
         >>> args = dict(model="rtdetr-l.yaml", data="coco8.yaml", imgsz=640, epochs=3)
         >>> trainer = RTDETRTrainer(overrides=args)
         >>> trainer.train()
+
+    Notes:
+        - F.grid_sample used in RT-DETR does not support the `deterministic=True` argument.
+        - AMP training can lead to NaN outputs and may produce errors during bipartite graph matching.
     """
 
     def get_model(self, cfg: dict | None = None, weights: str | None = None, verbose: bool = True):
-        """
-        Initialize and return an RT-DETR model for object detection tasks.
+        """Initialize and return an RT-DETR model for object detection tasks.
 
         Args:
             cfg (dict, optional): Model configuration.
@@ -60,8 +58,7 @@ class RTDETRTrainer(DetectionTrainer):
         return model
 
     def build_dataset(self, img_path: str, mode: str = "val", batch: int | None = None):
-        """
-        Build and return an RT-DETR dataset for training or validation.
+        """Build and return an RT-DETR dataset for training or validation.
 
         Args:
             img_path (str): Path to the folder containing images.