dgenerate-ultralytics-headless 8.3.143__py3-none-any.whl → 8.3.144__py3-none-any.whl

ultralytics/hub/session.py

@@ -5,6 +5,7 @@ import threading
 import time
 from http import HTTPStatus
 from pathlib import Path
+from typing import Any, Dict, Optional
 from urllib.parse import parse_qs, urlparse
 
 import requests
@@ -19,7 +20,7 @@ AGENT_NAME = f"python-{__version__}-colab" if IS_COLAB else f"python-{__version_
 
 class HUBTrainingSession:
     """
-    HUB training session for Ultralytics HUB YOLO models. Handles model initialization, heartbeats, and checkpointing.
+    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.
@@ -27,28 +28,29 @@ class HUBTrainingSession:
     Attributes:
         model_id (str): Identifier for the YOLO model being trained.
         model_url (str): URL for the model in Ultralytics HUB.
-        rate_limits (dict): Rate limits for different API calls (in seconds).
-        timers (dict): Timers for rate limiting.
-        metrics_queue (dict): Queue for the model's metrics.
-        metrics_upload_failed_queue (dict): Queue for metrics that failed to upload.
-        model (dict): Model data fetched from Ultralytics HUB.
+        rate_limits (Dict[str, int]): Rate limits for different API calls in seconds.
+        timers (Dict[str, Any]): Timers for rate limiting.
+        metrics_queue (Dict[str, Any]): Queue for the model's metrics.
+        metrics_upload_failed_queue (Dict[str, Any]): Queue for metrics that failed to upload.
+        model (Any): Model data fetched from Ultralytics HUB.
         model_file (str): Path to the model file.
-        train_args (dict): Arguments for training the model.
-        client (HUBClient): Client for interacting with Ultralytics HUB.
+        train_args (Dict[str, Any]): Arguments for training the model.
+        client (Any): Client for interacting with Ultralytics HUB.
         filename (str): Filename of the model.
 
     Examples:
+        Create a training session with a model URL
         >>> session = HUBTrainingSession("https://hub.ultralytics.com/models/example-model")
         >>> session.upload_metrics()
     """
 
-    def __init__(self, identifier):
+    def __init__(self, identifier: str):
         """
         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.
@@ -90,16 +92,16 @@ class HUBTrainingSession:
                 )
 
     @classmethod
-    def create_session(cls, identifier, args=None):
+    def create_session(cls, identifier: str, args: Optional[Dict[str, Any]] = None):
         """
         Create an authenticated HUBTrainingSession or return None.
 
         Args:
             identifier (str): Model identifier used to initialize the HUB training session.
-            args (dict, optional): Arguments for creating a new model if identifier is not a HUB model URL.
+            args (Dict[str, Any], optional): Arguments for creating a new model if identifier is not a HUB model URL.
 
         Returns:
-            (HUBTrainingSession | None): An authenticated session or None if creation fails.
+            session (HUBTrainingSession | None): An authenticated session or None if creation fails.
         """
         try:
             session = cls(identifier)
@@ -111,7 +113,7 @@ class HUBTrainingSession:
         except (PermissionError, ModuleNotFoundError, AssertionError):
             return None
 
-    def load_model(self, model_id):
+    def load_model(self, model_id: str):
         """
         Load an existing model from Ultralytics HUB using the provided model identifier.
 
@@ -137,12 +139,13 @@ class HUBTrainingSession:
         self.model.start_heartbeat(self.rate_limits["heartbeat"])
         LOGGER.info(f"{PREFIX}View model at {self.model_url} 🚀")
 
-    def create_model(self, model_args):
+    def create_model(self, model_args: Dict[str, Any]):
         """
         Initialize a HUB training session with the specified model arguments.
 
         Args:
-            model_args (dict): Arguments for creating the model, including batch size, epochs, image size, etc.
+            model_args (Dict[str, Any]): Arguments for creating the model, including batch size, epochs, image size,
+                etc.
 
         Returns:
             (None): If the model could not be created.
@@ -182,7 +185,7 @@ class HUBTrainingSession:
         LOGGER.info(f"{PREFIX}View model at {self.model_url} 🚀")
 
     @staticmethod
-    def _parse_identifier(identifier):
+    def _parse_identifier(identifier: str):
         """
         Parse the given identifier to determine the type and extract relevant components.
 
@@ -195,7 +198,9 @@ class HUBTrainingSession:
             identifier (str): The identifier string to be parsed.
 
         Returns:
-            (tuple): A tuple containing the API key, model ID, and filename as applicable.
+            api_key (str | None): Extracted API key if present.
+            model_id (str | None): Extracted model ID if present.
+            filename (str | None): Extracted filename if present.
 
         Raises:
             HUBModelError: If the identifier format is not recognized.
@@ -247,17 +252,17 @@ class HUBTrainingSession:
     def request_queue(
         self,
         request_func,
-        retry=3,
-        timeout=30,
-        thread=True,
-        verbose=True,
-        progress_total=None,
-        stream_response=None,
+        retry: int = 3,
+        timeout: int = 30,
+        thread: bool = True,
+        verbose: bool = True,
+        progress_total: Optional[int] = None,
+        stream_response: Optional[bool] = None,
         *args,
         **kwargs,
     ):
         """
-        Attempt to 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.
@@ -275,7 +280,7 @@ class HUBTrainingSession:
         """
 
         def retry_request():
-            """Attempt to call `request_func` with retries, timeout, and optional threading."""
+            """Attempt to call request_func with retries, timeout, and optional threading."""
             t0 = time.time()  # Record the start time for the timeout
             response = None
             for i in range(retry + 1):
@@ -327,16 +332,8 @@ class HUBTrainingSession:
             return retry_request()
 
     @staticmethod
-    def _should_retry(status_code):
-        """
-        Determine if a request should be retried based on the HTTP status code.
-
-        Args:
-            status_code (int): The HTTP status code from the response.
-
-        Returns:
-            (bool): True if the request should be retried, False otherwise.
-        """
+    def _should_retry(status_code: int) -> bool:
+        """Determine if a request should be retried based on the HTTP status code."""
         retry_codes = {
             HTTPStatus.REQUEST_TIMEOUT,
             HTTPStatus.BAD_GATEWAY,
@@ -344,7 +341,7 @@ class HUBTrainingSession:
         }
         return status_code in retry_codes
 
-    def _get_failure_message(self, response: requests.Response, retry: int, timeout: int):
+    def _get_failure_message(self, response: requests.Response, retry: int, timeout: int) -> str:
         """
         Generate a retry message based on the response status code.
 
@@ -423,24 +420,13 @@ class HUBTrainingSession:
 
     @staticmethod
     def _show_upload_progress(content_length: int, response: requests.Response) -> None:
-        """
-        Display a progress bar to track the upload progress of a file download.
-
-        Args:
-            content_length (int): The total size of the content to be downloaded in bytes.
-            response (requests.Response): The response object from the file download request.
-        """
+        """Display a progress bar to track the upload progress of a file download."""
         with TQDM(total=content_length, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
             for data in response.iter_content(chunk_size=1024):
                 pbar.update(len(data))
 
     @staticmethod
     def _iterate_content(response: requests.Response) -> None:
-        """
-        Process the streamed HTTP response data.
-
-        Args:
-            response (requests.Response): The response object from the file download request.
-        """
+        """Process the streamed HTTP response data."""
         for _ in response.iter_content(chunk_size=1024):
             pass  # Do nothing with data chunks

ultralytics/hub/utils.py

@@ -5,6 +5,7 @@ import random
 import threading
 import time
 from pathlib import Path
+from typing import Any, Optional
 
 import requests
 
@@ -36,7 +37,7 @@ PREFIX = colorstr("Ultralytics HUB: ")
 HELP_MSG = "If this issue persists please visit https://github.com/ultralytics/hub/issues for assistance."
 
 
-def request_with_credentials(url: str) -> any:
+def request_with_credentials(url: str) -> Any:
     """
     Make an AJAX request with cookies attached in a Google Colab environment.
 
@@ -77,7 +78,7 @@ def request_with_credentials(url: str) -> any:
     return output.eval_js("_hub_tmp")
 
 
-def requests_with_progress(method, url, **kwargs):
+def requests_with_progress(method: str, url: str, **kwargs) -> requests.Response:
     """
     Make an HTTP request using the specified method and URL, with an optional progress bar.
 
@@ -109,7 +110,17 @@ def requests_with_progress(method, url, **kwargs):
     return response
 
 
-def smart_request(method, url, retry=3, timeout=30, thread=True, code=-1, verbose=True, progress=False, **kwargs):
+def smart_request(
+    method: str,
+    url: str,
+    retry: int = 3,
+    timeout: int = 30,
+    thread: bool = True,
+    code: int = -1,
+    verbose: bool = True,
+    progress: bool = False,
+    **kwargs,
+) -> Optional[requests.Response]:
     """
     Make an HTTP request using the 'requests' library, with exponential backoff retries up to a specified timeout.
 
@@ -125,7 +136,8 @@ def smart_request(method, url, retry=3, timeout=30, thread=True, code=-1, verbos
         **kwargs (Any): Keyword arguments to be passed to the requests function specified in method.
 
     Returns:
-        (requests.Response): The HTTP response object. If the request is executed in a separate thread, returns None.
+        (requests.Response | None): The HTTP response object. If the request is executed in a separate thread, returns
+            None.
     """
     retry_codes = (408, 500)  # retry only these codes
 
@@ -177,7 +189,9 @@ class Events:
 
     Attributes:
         url (str): The URL to send anonymous events.
+        events (list): List of collected events to be sent.
         rate_limit (float): The rate limit in seconds for sending events.
+        t (float): Rate limit timer in seconds.
         metadata (dict): A dictionary containing metadata about the environment.
         enabled (bool): A flag to enable or disable Events based on certain conditions.
     """
@@ -214,7 +228,7 @@ class Events:
 
         Args:
             cfg (IterableSimpleNamespace): The configuration object containing mode and task information.
-            device (torch.device | str): The device type (e.g., 'cpu', 'cuda').
+            device (torch.device | str, optional): The device type (e.g., 'cpu', 'cuda').
         """
         if not self.enabled:
             # Events disabled, do nothing

ultralytics/models/fastsam/model.py

@@ -1,6 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from pathlib import Path
+from typing import Any, Dict, List, Optional
 
 from ultralytics.engine.model import Model
 
@@ -12,50 +13,67 @@ class FastSAM(Model):
     """
     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.
+    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.
 
     Attributes:
         model (str): Path to the pre-trained FastSAM model file.
         task (str): The task type, set to "segment" for FastSAM models.
 
+    Methods:
+        predict: Perform segmentation prediction on image or video source with optional prompts.
+        task_map: Returns mapping of segment task to predictor and validator classes.
+
     Examples:
+        Initialize FastSAM model and run prediction
         >>> from ultralytics import FastSAM
-        >>> model = FastSAM("last.pt")
+        >>> model = FastSAM("FastSAM-x.pt")
         >>> results = model.predict("ultralytics/assets/bus.jpg")
+
+        Run prediction with bounding box prompts
+        >>> results = model.predict("image.jpg", bboxes=[[100, 100, 200, 200]])
     """
 
-    def __init__(self, model="FastSAM-x.pt"):
+    def __init__(self, model: str = "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."
         super().__init__(model=model, task="segment")
 
-    def predict(self, source, stream=False, bboxes=None, points=None, labels=None, texts=None, **kwargs):
+    def predict(
+        self,
+        source,
+        stream: bool = False,
+        bboxes: Optional[List] = None,
+        points: Optional[List] = None,
+        labels: Optional[List] = None,
+        texts: Optional[List] = None,
+        **kwargs: Any,
+    ):
         """
         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.
+        prompts and passes them to the parent class predict method for processing.
 
         Args:
             source (str | PIL.Image | numpy.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): Bounding box coordinates for prompted segmentation in format [[x1, y1, x2, y2], ...].
-            points (list): Point coordinates for prompted segmentation in format [[x, y], ...].
-            labels (list): Class labels for prompted segmentation.
-            texts (list): Text prompts for segmentation guidance.
+            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]].
+            labels (List, optional): Class labels for prompted segmentation.
+            texts (List, optional): Text prompts for segmentation guidance.
             **kwargs (Any): Additional keyword arguments passed to the predictor.
 
         Returns:
-            (list): List of Results objects containing the prediction results.
+            (List): List of Results objects containing the prediction results.
         """
         prompts = dict(bboxes=bboxes, points=points, labels=labels, texts=texts)
         return super().predict(source, stream, prompts=prompts, **kwargs)
 
     @property
-    def task_map(self):
+    def task_map(self) -> Dict[str, Dict[str, Any]]:
         """Returns a dictionary mapping segment task to corresponding predictor and validator classes."""
         return {"segment": {"predictor": FastSAMPredictor, "validator": FastSAMValidator}}

ultralytics/models/fastsam/predict.py

@@ -26,10 +26,9 @@ 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):
@@ -41,7 +40,7 @@ class FastSAMPredictor(SegmentationPredictor):
         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.
         """
@@ -120,7 +119,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)

ultralytics/models/fastsam/utils.py

@@ -6,12 +6,12 @@ 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.
 
     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

@@ -6,9 +6,9 @@ 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.
 
@@ -18,6 +18,10 @@ class FastSAMValidator(SegmentationValidator):
         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):
@@ -25,11 +29,11 @@ class FastSAMValidator(SegmentationValidator):
         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.
+            pbar (tqdm.tqdm, optional): Progress bar for displaying progress.
+            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.

ultralytics/models/nas/model.py

@@ -9,6 +9,7 @@ Examples:
 """
 
 from pathlib import Path
+from typing import Any, Dict
 
 import torch
 
@@ -23,7 +24,7 @@ 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.
@@ -34,6 +35,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")
@@ -72,7 +76,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,7 +87,7 @@ 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):
+    def info(self, detailed: bool = False, verbose: bool = True) -> Dict[str, Any]:
         """
         Log model information.
 
@@ -92,11 +96,11 @@ class NAS(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

@@ -10,13 +10,13 @@ class NASPredictor(DetectionPredictor):
     """
     Ultralytics YOLO NAS Predictor for object detection.
 
-    This class extends the `DetectionPredictor` from Ultralytics engine and is responsible for post-processing the
+    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,7 +29,7 @@ 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):
@@ -53,6 +53,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

@@ -12,7 +12,7 @@ class NASValidator(DetectionValidator):
     """
     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 +25,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/model.py

@@ -21,13 +21,17 @@ class RTDETR(Model):
     """
     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.
 
+    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")
@@ -39,16 +43,13 @@ class RTDETR(Model):
 
         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'.
         """
         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

@@ -21,6 +21,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
@@ -37,14 +41,14 @@ class RTDETRPredictor(BasePredictor):
         model predictions to Results objects containing properly scaled bounding boxes.
 
         Args:
-            preds (List | Tuple): List of [predictions, extra] from the model, where predictions contain
+            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]
@@ -71,11 +75,14 @@ class RTDETRPredictor(BasePredictor):
 
     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.