ultralytics 8.3.89__py3-none-any.whl → 8.3.90__py3-none-any.whl

ultralytics/models/fastsam/model.py

@@ -10,7 +10,14 @@ from .val import FastSAMValidator
 
 class FastSAM(Model):
     """
-    FastSAM model interface.
+    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.
+
+    Attributes:
+        model (str): Path to the pre-trained FastSAM model file.
+        task (str): The task type, set to "segment" for FastSAM models.
 
     Examples:
         >>> from ultralytics import FastSAM
@@ -19,7 +26,7 @@ class FastSAM(Model):
     """
 
     def __init__(self, model="FastSAM-x.pt"):
-        """Call the __init__ method of the parent class (YOLO) with the updated default model."""
+        """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."
@@ -29,19 +36,21 @@ class FastSAM(Model):
         """
         Perform segmentation prediction on image or video source.
 
-        Supports prompted segmentation with bounding boxes, points, labels, and texts.
+        Supports prompted segmentation with bounding boxes, points, labels, and texts. The method packages these
+        prompts and passes them to the parent class predict method.
 
         Args:
-            source (str | PIL.Image | numpy.ndarray): Input source.
-            stream (bool): Enable real-time streaming.
-            bboxes (list): Bounding box coordinates for prompted segmentation.
-            points (list): Points for prompted segmentation.
-            labels (list): Labels for prompted segmentation.
-            texts (list): Texts for prompted segmentation.
-            **kwargs (Any): Additional keyword arguments.
+            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.
+            **kwargs (Any): Additional keyword arguments passed to the predictor.
 
         Returns:
-            (list): Model predictions.
+            (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)

ultralytics/models/fastsam/predict.py

@@ -13,21 +13,42 @@ from .utils import adjust_bboxes_to_image_border
 
 class FastSAMPredictor(SegmentationPredictor):
     """
-    FastSAMPredictor is specialized for fast SAM (Segment Anything Model) segmentation prediction tasks in Ultralytics
-    YOLO framework.
+    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-max suppression while optimizing for single-
-    class segmentation.
+    adjusts post-processing steps to incorporate mask prediction and non-maximum suppression while optimizing for
+    single-class segmentation.
+
+    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.
+
+    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.
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initializes a FastSAMPredictor for fast SAM segmentation tasks in Ultralytics YOLO framework."""
+        """Initialize the FastSAMPredictor with configuration and callbacks."""
         super().__init__(cfg, overrides, _callbacks)
         self.prompts = {}
 
     def postprocess(self, preds, img, orig_imgs):
-        """Applies box postprocess for FastSAM predictions."""
+        """
+        Apply postprocessing to FastSAM predictions and handle prompts.
+
+        Args:
+            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.
+
+        Returns:
+            (List[Results]): Processed results with prompts applied.
+        """
         bboxes = self.prompts.pop("bboxes", None)
         points = self.prompts.pop("points", None)
         labels = self.prompts.pop("labels", None)
@@ -46,18 +67,17 @@ class FastSAMPredictor(SegmentationPredictor):
 
     def prompt(self, results, bboxes=None, points=None, labels=None, texts=None):
         """
-        Internal function for image segmentation inference based on cues like bounding boxes, points, and masks.
-        Leverages SAM's specialized architecture for prompt-based, real-time segmentation.
+        Perform image segmentation inference based on cues like bounding boxes, points, and text prompts.
 
         Args:
-            results (Results | List[Results]): The original inference results from FastSAM models without any prompts.
+            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 contains string objects.
+            texts (str | List[str], optional): Textual prompts, a list containing string objects.
 
         Returns:
-            (List[Results]): The output results determined by 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
@@ -121,14 +141,14 @@ class FastSAMPredictor(SegmentationPredictor):
 
     def _clip_inference(self, images, texts):
         """
-        CLIP Inference process.
+        Perform CLIP inference to calculate similarity between images and text prompts.
 
         Args:
-            images (List[PIL.Image]): A list of source images and each of them should be PIL.Image type with RGB channel order.
-            texts (List[str]): A list of prompt texts and each of them should be 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): The similarity between given images and texts.
+            (torch.Tensor): Similarity matrix between given images and texts with shape (M, N).
         """
         try:
             import clip
@@ -146,5 +166,5 @@ class FastSAMPredictor(SegmentationPredictor):
         return (image_features * text_features[:, None]).sum(-1)  # (M, N)
 
     def set_prompts(self, prompts):
-        """Set prompts in advance."""
+        """Set prompts to be used during inference."""
         self.prompts = prompts

ultralytics/models/fastsam/utils.py

@@ -6,17 +6,17 @@ 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): (n, 4)
-        image_shape (tuple): (height, width)
-        threshold (int): pixel threshold
+        boxes (torch.Tensor): Bounding boxes with shape (n, 4) in xyxy format.
+        image_shape (Tuple[int, int]): Image dimensions as (height, width).
+        threshold (int): Pixel threshold for considering a box close to the border.
 
     Returns:
-        adjusted_boxes (torch.Tensor): adjusted bounding boxes
+       boxes (torch.Tensor): Adjusted bounding boxes with shape (n, 4).
     """
     # Image dimensions
     h, w = image_shape
 
-    # Adjust boxes
+    # Adjust boxes that are close to image borders
     boxes[boxes[:, 0] < threshold, 0] = 0  # x1
     boxes[boxes[:, 1] < threshold, 1] = 0  # y1
     boxes[boxes[:, 2] > w - threshold, 2] = w  # x2

ultralytics/models/fastsam/val.py

@@ -13,11 +13,11 @@ class FastSAMValidator(SegmentationValidator):
     to avoid errors during validation.
 
     Attributes:
-        dataloader: The data loader object used for validation.
-        save_dir (str): The directory where validation results will be saved.
-        pbar: A progress bar object.
-        args: Additional arguments for customization.
-        _callbacks: List of callback functions to be invoked during validation.
+        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.
     """
 
     def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
@@ -29,7 +29,7 @@ class FastSAMValidator(SegmentationValidator):
             save_dir (Path, optional): Directory to save results.
             pbar (tqdm.tqdm): Progress bar for displaying progress.
             args (SimpleNamespace): Configuration for the validator.
-            _callbacks (dict): Dictionary to store various callback functions.
+            _callbacks (List): 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

@@ -28,31 +28,39 @@ class NAS(Model):
     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.
+        task (str): The task type for the model, defaults to 'detect'.
+        predictor (NASPredictor): The predictor instance for making predictions.
+        validator (NASValidator): The validator instance for model validation.
+
     Examples:
         >>> from ultralytics import NAS
         >>> model = NAS("yolo_nas_s")
         >>> results = model.predict("ultralytics/assets/bus.jpg")
 
-    Attributes:
-        model (str): Path to the pre-trained model or model name. Defaults to 'yolo_nas_s.pt'.
-
-    Note:
+    Notes:
         YOLO-NAS models only support pre-trained models. Do not provide YAML configuration files.
     """
 
-    def __init__(self, model="yolo_nas_s.pt") -> None:
-        """Initializes the NAS model with the provided or default 'yolo_nas_s.pt' model."""
+    def __init__(self, model: str = "yolo_nas_s.pt") -> None:
+        """Initialize the NAS model with the provided or default model."""
         assert Path(model).suffix not in {".yaml", ".yml"}, "YOLO-NAS models only support pre-trained models."
         super().__init__(model, task="detect")
 
     def _load(self, weights: str, task=None) -> None:
-        """Loads an existing NAS model weights or creates a new NAS model with pretrained weights if not provided."""
+        """
+        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.
+            task (str, optional): Task type for the model.
+        """
         import super_gradients
 
         suffix = Path(weights).suffix
         if suffix == ".pt":
             self.model = torch.load(attempt_download_asset(weights))
-
         elif suffix == "":
             self.model = super_gradients.training.models.get(weights, pretrained_weights="coco")
 
@@ -74,17 +82,20 @@ class NAS(Model):
         self.model.task = "detect"  # for export()
         self.model.args = {**DEFAULT_CFG_DICT, **self.overrides}  # for export()
 
-    def info(self, detailed=False, verbose=True):
+    def info(self, detailed: bool = False, verbose: bool = True):
         """
-        Logs model info.
+        Log model information.
 
         Args:
             detailed (bool): Show detailed information about model.
             verbose (bool): Controls verbosity.
+
+        Returns:
+            (dict): Model information dictionary.
         """
         return model_info(self.model, detailed=detailed, verbose=verbose, imgsz=640)
 
     @property
     def task_map(self):
-        """Returns a dictionary mapping tasks to respective predictor and validator classes."""
+        """Return a dictionary mapping tasks to respective predictor and validator classes."""
         return {"detect": {"predictor": NASPredictor, "validator": NASValidator}}

ultralytics/models/nas/predict.py

@@ -16,26 +16,30 @@ class NASPredictor(BasePredictor):
     scaling the bounding boxes to fit the original image dimensions.
 
     Attributes:
-        args (Namespace): Namespace containing various configurations for post-processing.
+        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.
 
     Examples:
         >>> from ultralytics import NAS
         >>> model = NAS("yolo_nas_s")
         >>> predictor = model.predictor
 
-        Assumes that raw_preds, img, orig_imgs are available
+        Assume that raw_preds, img, orig_imgs are available
         >>> results = predictor.postprocess(raw_preds, img, orig_imgs)
 
-    Note:
+    Notes:
         Typically, this class is not instantiated directly. It is used internally within the `NAS` class.
     """
 
     def postprocess(self, preds_in, img, orig_imgs):
         """Postprocess predictions and returns a list of Results objects."""
-        # Cat boxes and class scores
+        # Convert boxes from xyxy to xywh format and concatenate with class scores
         boxes = ops.xyxy2xywh(preds_in[0][0])
         preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)
 
+        # Apply non-maximum suppression to filter overlapping detections
         preds = ops.non_max_suppression(
             preds,
             self.args.conf,
@@ -50,6 +54,7 @@ class NASPredictor(BasePredictor):
 
         results = []
         for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
+            # Scale bounding boxes to match original image dimensions
             pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
             results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
         return results

ultralytics/models/nas/val.py

@@ -17,25 +17,25 @@ class NASValidator(DetectionValidator):
     ultimately producing the final detections.
 
     Attributes:
-        args (Namespace): Namespace containing various configurations for post-processing, such as confidence and IoU.
+        args (Namespace): Namespace containing various configurations for post-processing, such as confidence and IoU
+            thresholds.
         lb (torch.Tensor): Optional tensor for multilabel NMS.
 
     Examples:
         >>> from ultralytics import NAS
         >>> model = NAS("yolo_nas_s")
         >>> validator = model.validator
-
         Assumes that raw_preds are available
         >>> final_preds = validator.postprocess(raw_preds)
 
-    Note:
+    Notes:
         This class is generally not instantiated directly but is used internally within the `NAS` class.
     """
 
     def postprocess(self, preds_in):
         """Apply Non-maximum suppression to prediction outputs."""
-        boxes = ops.xyxy2xywh(preds_in[0][0])
-        preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)
+        boxes = ops.xyxy2xywh(preds_in[0][0])  # Convert bounding box format from xyxy to xywh
+        preds = torch.cat((boxes, preds_in[0][1]), -1).permute(0, 2, 1)  # Concatenate boxes with scores and permute
         return super().postprocess(
             preds,
             max_time_img=0.5,

ultralytics/models/rtdetr/model.py

@@ -1,10 +1,12 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 """
-Interface for Baidu's RT-DETR, a Vision Transformer-based real-time object detector. RT-DETR offers real-time
-performance and high accuracy, excelling in accelerated backends like CUDA with TensorRT. It features an efficient
-hybrid encoder and IoU-aware query selection for enhanced detection accuracy.
+Interface for Baidu's RT-DETR, a Vision Transformer-based real-time object detector.
 
-For more information on RT-DETR, visit: https://arxiv.org/pdf/2304.08069.pdf
+RT-DETR offers real-time performance and high accuracy, excelling in accelerated backends like CUDA with TensorRT.
+It features an efficient hybrid encoder and IoU-aware query selection for enhanced detection accuracy.
+
+References:
+    https://arxiv.org/pdf/2304.08069.pdf
 """
 
 from ultralytics.engine.model import Model
@@ -17,19 +19,26 @@ from .val import RTDETRValidator
 
 class RTDETR(Model):
     """
-    Interface for Baidu's RT-DETR model. This Vision Transformer-based object detector provides real-time performance
-    with high accuracy. It supports efficient hybrid encoding, IoU-aware query selection, and adaptable inference speed.
+    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.
 
     Attributes:
-        model (str): Path to the pre-trained model. Defaults to 'rtdetr-l.pt'.
+        model (str): Path to the pre-trained model.
+
+    Examples:
+        >>> from ultralytics import RTDETR
+        >>> model = RTDETR("rtdetr-l.pt")
+        >>> results = model("image.jpg")
     """
 
-    def __init__(self, model="rtdetr-l.pt") -> None:
+    def __init__(self, model: str = "rtdetr-l.pt") -> None:
         """
-        Initializes the RT-DETR model with the given pre-trained model file. Supports .pt and .yaml formats.
+        Initialize the RT-DETR model with the given pre-trained model file.
 
         Args:
-            model (str): Path to the pre-trained model. Defaults to 'rtdetr-l.pt'.
+            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'.
@@ -42,7 +51,7 @@ class RTDETR(Model):
         Returns 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.
+            (Dict): A dictionary mapping task names to Ultralytics task classes for the RT-DETR model.
         """
         return {
             "detect": {

ultralytics/models/rtdetr/predict.py

@@ -10,11 +10,16 @@ from ultralytics.utils import ops
 
 class RTDETRPredictor(BasePredictor):
     """
-    RT-DETR (Real-Time Detection Transformer) Predictor extending the BasePredictor class for making predictions using
-    Baidu's RT-DETR model.
+    RT-DETR (Real-Time Detection Transformer) Predictor extending the BasePredictor class for making predictions.
 
-    This class leverages the power of Vision Transformers to provide real-time object detection while maintaining
-    high accuracy. It supports key features like efficient hybrid encoding and IoU-aware query selection.
+    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).
+        args (dict): Argument overrides for the predictor.
+        model (torch.nn.Module): The loaded RT-DETR model.
+        batch (List): Current batch of processed inputs.
 
     Examples:
         >>> from ultralytics.utils import ASSETS
@@ -22,25 +27,23 @@ class RTDETRPredictor(BasePredictor):
         >>> args = dict(model="rtdetr-l.pt", source=ASSETS)
         >>> predictor = RTDETRPredictor(overrides=args)
         >>> predictor.predict_cli()
-
-    Attributes:
-        imgsz (int): Image size for inference (must be square and scale-filled).
-        args (dict): Argument overrides for the predictor.
     """
 
     def postprocess(self, preds, img, orig_imgs):
         """
         Postprocess the raw predictions from the model to generate bounding boxes and confidence scores.
 
-        The method filters detections based on confidence and class if specified in `self.args`.
+        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): List of [predictions, extra] from the model.
-            img (torch.Tensor): Processed input images.
-            orig_imgs (list or torch.Tensor): Original, unprocessed images.
+            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:
-            (list[Results]): A list of Results objects containing the post-processed bounding boxes, confidence scores,
+            (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
@@ -61,8 +64,8 @@ class RTDETRPredictor(BasePredictor):
                 idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx
             pred = torch.cat([bbox, max_score, cls], dim=-1)[idx]  # filter
             oh, ow = orig_img.shape[:2]
-            pred[..., [0, 2]] *= ow
-            pred[..., [1, 3]] *= oh
+            pred[..., [0, 2]] *= ow  # scale x coordinates to original width
+            pred[..., [1, 3]] *= oh  # scale y coordinates to original height
             results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
         return results
 

ultralytics/models/rtdetr/train.py

@@ -13,9 +13,18 @@ from .val import RTDETRDataset, RTDETRValidator
 
 class RTDETRTrainer(DetectionTrainer):
     """
-    Trainer class for the RT-DETR model developed by Baidu for real-time object detection. Extends the DetectionTrainer
-    class for YOLO to adapt to the specific features and architecture of RT-DETR. This model leverages Vision
-    Transformers and has capabilities like IoU-aware query selection and adaptable inference speed.
+    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.
+
+    Attributes:
+        loss_names (Tuple[str]): 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.
@@ -33,9 +42,9 @@ class RTDETRTrainer(DetectionTrainer):
         Initialize and return an RT-DETR model for object detection tasks.
 
         Args:
-            cfg (dict, optional): Model configuration. Defaults to None.
-            weights (str, optional): Path to pre-trained model weights. Defaults to None.
-            verbose (bool): Verbose logging if True. Defaults to True.
+            cfg (Dict, optional): Model configuration.
+            weights (str, optional): Path to pre-trained model weights.
+            verbose (bool): Verbose logging if True.
 
         Returns:
             (RTDETRDetectionModel): Initialized model.
@@ -52,7 +61,7 @@ class RTDETRTrainer(DetectionTrainer):
         Args:
             img_path (str): Path to the folder containing images.
             mode (str): Dataset mode, either 'train' or 'val'.
-            batch (int, optional): Batch size for rectangle training. Defaults to None.
+            batch (int, optional): Batch size for rectangle training.
 
         Returns:
             (RTDETRDataset): Dataset object for the specific mode.
@@ -73,24 +82,19 @@ class RTDETRTrainer(DetectionTrainer):
         )
 
     def get_validator(self):
-        """
-        Returns a DetectionValidator suitable for RT-DETR model validation.
-
-        Returns:
-            (RTDETRValidator): Validator object for model validation.
-        """
+        """Returns 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))
 
     def preprocess_batch(self, batch):
         """
-        Preprocess a batch of images. Scales and converts the images to float format.
+        Preprocess a batch of images by scaling and converting to float format.
 
         Args:
-            batch (dict): Dictionary containing a batch of images, bboxes, and labels.
+            batch (Dict): Dictionary containing a batch of images, bboxes, and labels.
 
         Returns:
-            (dict): Preprocessed batch.
+            (Dict): Preprocessed batch with ground truth bounding boxes and classes separated by batch index.
         """
         batch = super().preprocess_batch(batch)
         bs = len(batch["img"])

ultralytics/models/rtdetr/val.py

@@ -22,13 +22,20 @@ class RTDETRDataset(YOLODataset):
         """Initialize the RTDETRDataset class by inheriting from the YOLODataset class."""
         super().__init__(*args, data=data, **kwargs)
 
-    # NOTE: add stretch version load_image for RTDETR mosaic
     def load_image(self, i, rect_mode=False):
         """Loads 1 image from dataset index 'i', returns (im, resized hw)."""
         return super().load_image(i=i, rect_mode=rect_mode)
 
     def build_transforms(self, hyp=None):
-        """Temporary, only for evaluation."""
+        """
+        Build transformation pipeline for the dataset.
+
+        Args:
+            hyp (Dict, optional): Hyperparameters for transformations.
+
+        Returns:
+            (Compose): Composition of transformation functions.
+        """
         if self.augment:
             hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0
             hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0
@@ -75,7 +82,10 @@ class RTDETRValidator(DetectionValidator):
         Args:
             img_path (str): Path to the folder containing images.
             mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode.
-            batch (int, optional): Size of batches, this is for `rect`. Defaults to None.
+            batch (int, optional): Size of batches, this is for `rect`.
+
+        Returns:
+            (RTDETRDataset): Dataset configured for RT-DETR validation.
         """
         return RTDETRDataset(
             img_path=img_path,
@@ -90,7 +100,15 @@ class RTDETRValidator(DetectionValidator):
         )
 
     def postprocess(self, preds):
-        """Apply Non-maximum suppression to prediction outputs."""
+        """
+        Apply Non-maximum suppression to prediction outputs.
+
+        Args:
+            preds (List | Tuple | torch.Tensor): Raw predictions from the model.
+
+        Returns:
+            (List[torch.Tensor]): List of processed predictions for each image in batch.
+        """
         if not isinstance(preds, (list, tuple)):  # list for PyTorch inference but list[0] Tensor for export inference
             preds = [preds, None]
 
@@ -111,7 +129,16 @@ class RTDETRValidator(DetectionValidator):
         return outputs
 
     def _prepare_batch(self, si, batch):
-        """Prepares a batch for training or inference by applying transformations."""
+        """
+        Prepares a batch for validation by applying necessary transformations.
+
+        Args:
+            si (int): Batch index.
+            batch (Dict): Batch data containing images and annotations.
+
+        Returns:
+            (Dict): Prepared batch with transformed annotations.
+        """
         idx = batch["batch_idx"] == si
         cls = batch["cls"][idx].squeeze(-1)
         bbox = batch["bboxes"][idx]
@@ -125,7 +152,16 @@ class RTDETRValidator(DetectionValidator):
         return {"cls": cls, "bbox": bbox, "ori_shape": ori_shape, "imgsz": imgsz, "ratio_pad": ratio_pad}
 
     def _prepare_pred(self, pred, pbatch):
-        """Prepares and returns a batch with transformed bounding boxes and class labels."""
+        """
+        Prepares predictions by scaling bounding boxes to original image dimensions.
+
+        Args:
+            pred (torch.Tensor): Raw predictions.
+            pbatch (Dict): Prepared batch information.
+
+        Returns:
+            (torch.Tensor): Predictions scaled to original image dimensions.
+        """
         predn = pred.clone()
         predn[..., [0, 2]] *= pbatch["ori_shape"][1] / self.args.imgsz  # native-space pred
         predn[..., [1, 3]] *= pbatch["ori_shape"][0] / self.args.imgsz  # native-space pred

ultralytics/models/sam/__init__.py

@@ -3,4 +3,4 @@
 from .model import SAM
 from .predict import Predictor, SAM2Predictor, SAM2VideoPredictor
 
-__all__ = "SAM", "Predictor", "SAM2Predictor", "SAM2VideoPredictor"  # tuple or list
+__all__ = "SAM", "Predictor", "SAM2Predictor", "SAM2VideoPredictor"  # tuple or list of exportable items