dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

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
@@ -26,22 +26,20 @@ class FastSAMPredictor(SegmentationPredictor):
         clip_preprocess (Any, optional): CLIP preprocessing function for images, loaded on demand.
 
     Methods:
-        postprocess: Applies box postprocessing for FastSAM predictions.
-        prompt: Performs image segmentation inference based on various prompt types.
-        _clip_inference: Performs CLIP inference to calculate similarity between images and text prompts.
-        set_prompts: Sets prompts to be used during inference.
+        postprocess: Apply postprocessing to FastSAM predictions and handle prompts.
+        prompt: Perform image segmentation inference based on various prompt types.
+        set_prompts: Set prompts to be used during inference.
     """
 
     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
         optimized for single-class segmentation.
 
         Args:
-            cfg (dict): Configuration for the predictor. Defaults to Ultralytics DEFAULT_CFG.
+            cfg (dict): Configuration for the predictor.
             overrides (dict, optional): Configuration overrides.
             _callbacks (list, optional): List of callback functions.
         """
@@ -49,16 +47,15 @@ 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.
+            preds (list[torch.Tensor]): Raw predictions from the model.
             img (torch.Tensor): Input image tensor that was fed to the model.
-            orig_imgs (List[numpy.ndarray]): Original images before preprocessing.
+            orig_imgs (list[np.ndarray]): Original images before preprocessing.
 
         Returns:
-            (List[Results]): Processed results with prompts applied.
+            (list[Results]): Processed results with prompts applied.
         """
         bboxes = self.prompts.pop("bboxes", None)
         points = self.prompts.pop("points", None)
@@ -77,18 +74,17 @@ 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.
-            bboxes (np.ndarray | List, optional): Bounding boxes with shape (N, 4), in XYXY format.
-            points (np.ndarray | List, optional): Points indicating object locations with shape (N, 2), in pixels.
-            labels (np.ndarray | List, optional): Labels for point prompts, shape (N, ). 1 = foreground, 0 = background.
-            texts (str | List[str], optional): Textual prompts, a list containing string objects.
+            results (Results | list[Results]): Original inference results from FastSAM models without any prompts.
+            bboxes (np.ndarray | list, optional): Bounding boxes with shape (N, 4), in XYXY format.
+            points (np.ndarray | list, optional): Points indicating object locations with shape (N, 2), in pixels.
+            labels (np.ndarray | list, optional): Labels for point prompts, shape (N, ). 1 = foreground, 0 = background.
+            texts (str | list[str], optional): Textual prompts, a list containing string objects.
 
         Returns:
-            (List[Results]): Output results filtered and determined by the provided prompts.
+            (list[Results]): Output results filtered and determined by the provided prompts.
         """
         if bboxes is None and points is None and texts is None:
             return results
@@ -101,7 +97,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:
@@ -120,7 +116,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"Excepted `labels` got same size as `point`, but got {len(labels)} and {len(points)}"
+                    f"Expected `labels` with same size as `point`, but got {len(labels)} and {len(points)}"
                 )
                 point_idx = (
                     torch.ones(len(result), dtype=torch.bool, device=self.device)
@@ -136,7 +132,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]))
@@ -151,30 +147,23 @@ 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.
-            texts (List[str]): List of prompt texts, each should be a string object.
+            images (list[PIL.Image]): List of source images, each should be PIL.Image with RGB channel order.
+            texts (list[str]): List of prompt texts, each should be a string object.
 
         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,16 +2,15 @@
 
 
 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.
-        image_shape (Tuple[int, int]): Image dimensions as (height, width).
+        boxes (torch.Tensor): Bounding boxes with shape (N, 4) in xyxy format.
+        image_shape (tuple): Image dimensions as (height, width).
         threshold (int): Pixel threshold for considering a box close to the border.
 
     Returns:
-       boxes (torch.Tensor): Adjusted bounding boxes with shape (n, 4).
+        (torch.Tensor): Adjusted bounding boxes with shape (N, 4).
     """
     # Image dimensions
     h, w = image_shape

ultralytics/models/fastsam/val.py

@@ -1,40 +1,38 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from ultralytics.models.yolo.segment import SegmentationValidator
-from ultralytics.utils.metrics import SegmentMetrics
 
 
 class FastSAMValidator(SegmentationValidator):
-    """
-    Custom validation class for fast SAM (Segment Anything Model) segmentation in Ultralytics YOLO framework.
+    """Custom validation class for Fast SAM (Segment Anything Model) segmentation in 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 Fast SAM. This class
     sets the task to 'segment' and uses the SegmentMetrics for evaluation. Additionally, plotting features are disabled
     to avoid errors during validation.
 
     Attributes:
         dataloader (torch.utils.data.DataLoader): The data loader object used for validation.
         save_dir (Path): The directory where validation results will be saved.
-        pbar (tqdm.tqdm): A progress bar object for displaying validation progress.
         args (SimpleNamespace): Additional arguments for customization of the validation process.
         _callbacks (list): List of callback functions to be invoked during validation.
+        metrics (SegmentMetrics): Segmentation metrics calculator for evaluation.
+
+    Methods:
+        __init__: Initialize the FastSAMValidator with custom settings for Fast SAM.
     """
 
-    def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
-        """
-        Initialize the FastSAMValidator class, setting the task to 'segment' and metrics to SegmentMetrics.
+    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.
 
         Args:
-            dataloader (torch.utils.data.DataLoader): 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.
-            pbar (tqdm.tqdm): Progress bar for displaying progress.
-            args (SimpleNamespace): Configuration for the validator.
-            _callbacks (list): List of callback functions to be invoked during validation.
+            args (SimpleNamespace, optional): Configuration for the validator.
+            _callbacks (list, optional): List of callback functions to be invoked during validation.
 
         Notes:
             Plots for ConfusionMatrix and other related metrics are disabled in this class to avoid errors.
         """
-        super().__init__(dataloader, save_dir, pbar, args, _callbacks)
+        super().__init__(dataloader, save_dir, args, _callbacks)
         self.args.task = "segment"
         self.args.plots = False  # disable ConfusionMatrix and other plots to avoid errors
-        self.metrics = SegmentMetrics(save_dir=self.save_dir)

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

@@ -1,20 +1,16 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
-"""
-YOLO-NAS model interface.
 
-Examples:
-    >>> from ultralytics import NAS
-    >>> model = NAS("yolo_nas_s")
-    >>> results = model.predict("ultralytics/assets/bus.jpg")
-"""
+from __future__ import annotations
 
 from pathlib import Path
+from typing import Any
 
 import torch
 
 from ultralytics.engine.model import Model
 from ultralytics.utils import DEFAULT_CFG_DICT
 from ultralytics.utils.downloads import attempt_download_asset
+from ultralytics.utils.patches import torch_load
 from ultralytics.utils.torch_utils import model_info
 
 from .predict import NASPredictor
@@ -22,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.
@@ -34,6 +29,9 @@ class NAS(Model):
         predictor (NASPredictor): The predictor instance for making predictions.
         validator (NASValidator): The validator instance for model validation.
 
+    Methods:
+        info: Log model information and return model details.
+
     Examples:
         >>> from ultralytics import NAS
         >>> model = NAS("yolo_nas_s")
@@ -49,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.
@@ -60,7 +57,7 @@ class NAS(Model):
 
         suffix = Path(weights).suffix
         if suffix == ".pt":
-            self.model = torch.load(attempt_download_asset(weights))
+            self.model = torch_load(attempt_download_asset(weights))
         elif suffix == "":
             self.model = super_gradients.training.models.get(weights, pretrained_weights="coco")
 
@@ -72,7 +69,7 @@ class NAS(Model):
         self.model._original_forward = self.model.forward
         self.model.forward = new_forward
 
-        # Standardize model
+        # Standardize model attributes for compatibility
         self.model.fuse = lambda verbose=True: self.model
         self.model.stride = torch.tensor([32])
         self.model.names = dict(enumerate(self.model._class_names))
@@ -83,20 +80,19 @@ class NAS(Model):
         self.model.args = {**DEFAULT_CFG_DICT, **self.overrides}  # for export()
         self.model.eval()
 
-    def info(self, detailed: bool = False, verbose: bool = True):
-        """
-        Log model information.
+    def info(self, detailed: bool = False, verbose: bool = True) -> dict[str, Any]:
+        """Log model information.
 
         Args:
             detailed (bool): Show detailed information about model.
             verbose (bool): Controls verbosity.
 
         Returns:
-            (dict): Model information dictionary.
+            (dict[str, Any]): Model information dictionary.
         """
         return model_info(self.model, detailed=detailed, verbose=verbose, imgsz=640)
 
     @property
-    def task_map(self):
+    def task_map(self) -> dict[str, dict[str, Any]]:
         """Return a dictionary mapping tasks to respective predictor and validator classes."""
         return {"detect": {"predictor": NASPredictor, "validator": NASValidator}}

ultralytics/models/nas/predict.py

@@ -7,16 +7,15 @@ 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 threshold,
-            IoU threshold, agnostic NMS flag, maximum detections, and class filtering options.
+        args (Namespace): Namespace containing various configurations for post-processing including confidence
+            threshold, IoU threshold, agnostic NMS flag, maximum detections, and class filtering options.
         model (torch.nn.Module): The YOLO NAS model used for inference.
         batch (list): Batch of inputs for processing.
 
@@ -29,16 +28,15 @@ class NASPredictor(DetectionPredictor):
         >>> results = predictor.postprocess(raw_preds, img, orig_imgs)
 
     Notes:
-        Typically, this class is not instantiated directly. It is used internally within the `NAS` class.
+        Typically, this class is not instantiated directly. It is used internally within the NAS class.
     """
 
     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.
@@ -53,6 +51,6 @@ class NASPredictor(DetectionPredictor):
             >>> predictor = NAS("yolo_nas_s").predictor
             >>> results = predictor.postprocess(raw_preds, img, orig_imgs)
         """
-        boxes = ops.xyxy2xywh(preds_in[0][0])
-        preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)  # concatenate with class scores
+        boxes = ops.xyxy2xywh(preds_in[0][0])  # Convert bounding boxes from xyxy to xywh format
+        preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)  # Concatenate boxes with class scores
         return super().postprocess(preds, img, orig_imgs)

ultralytics/models/nas/val.py

@@ -9,10 +9,9 @@ __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
+    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,
     ultimately producing the final detections.
 
@@ -25,11 +24,11 @@ class NASValidator(DetectionValidator):
         >>> from ultralytics import NAS
         >>> model = NAS("yolo_nas_s")
         >>> validator = model.validator
-        Assumes that raw_preds are available
+        >>> # Assumes that raw_preds are available
         >>> final_preds = validator.postprocess(raw_preds)
 
     Notes:
-        This class is generally not instantiated directly but is used internally within the `NAS` class.
+        This class is generally not instantiated directly but is used internally within the NAS class.
     """
 
     def postprocess(self, preds_in):

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,8 +19,7 @@ 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.
@@ -27,28 +27,28 @@ class RTDETR(Model):
     Attributes:
         model (str): Path to the pre-trained model.
 
+    Methods:
+        task_map: Return a task map for RT-DETR, associating tasks with corresponding Ultralytics classes.
+
     Examples:
+        Initialize RT-DETR with a pre-trained model
         >>> from ultralytics import RTDETR
         >>> model = RTDETR("rtdetr-l.pt")
         >>> results = model("image.jpg")
     """
 
     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.
-
-        Raises:
-            NotImplementedError: If the model file extension is not 'pt', 'yaml', or 'yml'.
         """
+        assert TORCH_1_11, "RTDETR requires torch>=1.11"
         super().__init__(model=model, task="detect")
 
     @property
     def task_map(self) -> dict:
-        """
-        Returns 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).
@@ -21,6 +20,10 @@ class RTDETRPredictor(BasePredictor):
         model (torch.nn.Module): The loaded RT-DETR model.
         batch (list): Current batch of processed inputs.
 
+    Methods:
+        postprocess: Postprocess raw model predictions to generate bounding boxes and confidence scores.
+        pre_transform: Pre-transform input images before feeding them into the model for inference.
+
     Examples:
         >>> from ultralytics.utils import ASSETS
         >>> from ultralytics.models.rtdetr import RTDETRPredictor
@@ -30,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.
+            orig_imgs (list | torch.Tensor): Original, unprocessed images.
 
         Returns:
-            (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]
@@ -63,6 +65,7 @@ class RTDETRPredictor(BasePredictor):
             if self.args.classes is not None:
                 idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx
             pred = torch.cat([bbox, max_score, cls], dim=-1)[idx]  # filter
+            pred = pred[pred[:, 4].argsort(descending=True)][: self.args.max_det]
             oh, ow = orig_img.shape[:2]
             pred[..., [0, 2]] *= ow  # scale x coordinates to original width
             pred[..., [1, 3]] *= oh  # scale y coordinates to original height
@@ -70,12 +73,14 @@ class RTDETRPredictor(BasePredictor):
         return results
 
     def pre_transform(self, 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.
+        """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.
 
         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

@@ -1,5 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from copy import copy
 
 from ultralytics.models.yolo.detect import DetectionTrainer
@@ -10,34 +12,37 @@ 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[str]): Names of the loss components used for training.
+        loss_names (tuple): Names of the loss components used for training.
         data (dict): Dataset configuration containing class count and other parameters.
         args (dict): Training arguments and hyperparameters.
         save_dir (Path): Directory to save training results.
         test_loader (DataLoader): DataLoader for validation/testing data.
 
-    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.
+    Methods:
+        get_model: Initialize and return an RT-DETR model for object detection tasks.
+        build_dataset: Build and return an RT-DETR dataset for training or validation.
+        get_validator: Return a DetectionValidator suitable for RT-DETR model validation.
 
     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=None, weights=None, verbose=True):
-        """
-        Initialize and return an RT-DETR model for object detection tasks.
+    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.
 
         Args:
             cfg (dict, optional): Model configuration.
@@ -52,9 +57,8 @@ class RTDETRTrainer(DetectionTrainer):
             model.load(weights)
         return model
 
-    def build_dataset(self, img_path, mode="val", batch=None):
-        """
-        Build and return an RT-DETR dataset for training or validation.
+    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.
 
         Args:
             img_path (str): Path to the folder containing images.
@@ -80,6 +84,6 @@ class RTDETRTrainer(DetectionTrainer):
         )
 
     def get_validator(self):
-        """Returns a DetectionValidator suitable for RT-DETR model validation."""
+        """Return a DetectionValidator suitable for RT-DETR model validation."""
         self.loss_names = "giou_loss", "cls_loss", "l1_loss"
         return RTDETRValidator(self.test_loader, save_dir=self.save_dir, args=copy(self.args))