ultralytics 8.3.100__py3-none-any.whl → 8.3.102__py3-none-any.whl

ultralytics/models/sam/modules/memory_attention.py

@@ -60,7 +60,17 @@ class MemoryAttentionLayer(nn.Module):
         pos_enc_at_cross_attn_keys: bool = True,
         pos_enc_at_cross_attn_queries: bool = False,
     ):
-        """Initialize a memory attention layer with self-attention, cross-attention, and feedforward components."""
+        """
+        Initialize a memory attention layer with self-attention, cross-attention, and feedforward components.
+
+        Args:
+            d_model (int): Dimensionality of the model.
+            dim_feedforward (int): Dimensionality of the feedforward network.
+            dropout (float): Dropout rate for regularization.
+            pos_enc_at_attn (bool): Whether to add positional encoding at attention.
+            pos_enc_at_cross_attn_keys (bool): Whether to add positional encoding to cross-attention keys.
+            pos_enc_at_cross_attn_queries (bool): Whether to add positional encoding to cross-attention queries.
+        """
         super().__init__()
         self.d_model = d_model
         self.dim_feedforward = dim_feedforward
@@ -183,7 +193,31 @@ class MemoryAttention(nn.Module):
         num_layers: int,
         batch_first: bool = True,  # Do layers expect batch first input?
     ):
-        """Initialize MemoryAttention with specified layers and normalization for sequential data processing."""
+        """
+        Initialize MemoryAttention with specified layers and normalization for sequential data processing.
+
+        This class implements a multi-layer attention mechanism that combines self-attention and cross-attention
+        for processing sequential data, particularly useful in transformer-like architectures.
+
+        Args:
+            d_model (int): The dimension of the model's hidden state.
+            pos_enc_at_input (bool): Whether to apply positional encoding at the input.
+            layer (nn.Module): The attention layer to be used in the module.
+            num_layers (int): The number of attention layers.
+            batch_first (bool): Whether the input tensors are in batch-first format.
+
+        Examples:
+            >>> d_model = 256
+            >>> layer = MemoryAttentionLayer(d_model)
+            >>> attention = MemoryAttention(d_model, pos_enc_at_input=True, layer=layer, num_layers=3)
+            >>> curr = torch.randn(10, 32, d_model)  # (seq_len, batch_size, d_model)
+            >>> memory = torch.randn(20, 32, d_model)  # (mem_len, batch_size, d_model)
+            >>> curr_pos = torch.randn(10, 32, d_model)
+            >>> memory_pos = torch.randn(20, 32, d_model)
+            >>> output = attention(curr, memory, curr_pos, memory_pos)
+            >>> print(output.shape)
+            torch.Size([10, 32, 256])
+        """
         super().__init__()
         self.d_model = d_model
         self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(num_layers)])
@@ -200,7 +234,31 @@ class MemoryAttention(nn.Module):
         memory_pos: Optional[Tensor] = None,  # pos_enc for cross-attention inputs
         num_obj_ptr_tokens: int = 0,  # number of object pointer *tokens*
     ) -> torch.Tensor:
-        """Process inputs through attention layers, applying self and cross-attention with positional encoding."""
+        """
+        Process inputs through attention layers, applying self and cross-attention with positional encoding.
+
+        Args:
+            curr (torch.Tensor): Self-attention input tensor, representing the current state.
+            memory (torch.Tensor): Cross-attention input tensor, representing memory information.
+            curr_pos (Optional[Tensor]): Positional encoding for self-attention inputs.
+            memory_pos (Optional[Tensor]): Positional encoding for cross-attention inputs.
+            num_obj_ptr_tokens (int): Number of object pointer tokens to exclude from rotary position embedding.
+
+        Returns:
+            (torch.Tensor): Processed output tensor after applying attention layers and normalization.
+
+        Examples:
+            >>> d_model = 256
+            >>> layer = MemoryAttentionLayer(d_model)
+            >>> attention = MemoryAttention(d_model, pos_enc_at_input=True, layer=layer, num_layers=3)
+            >>> curr = torch.randn(10, 32, d_model)  # (seq_len, batch_size, d_model)
+            >>> memory = torch.randn(20, 32, d_model)  # (mem_len, batch_size, d_model)
+            >>> curr_pos = torch.randn(10, 32, d_model)
+            >>> memory_pos = torch.randn(20, 32, d_model)
+            >>> output = attention(curr, memory, curr_pos, memory_pos)
+            >>> print(output.shape)
+            torch.Size([10, 32, 256])
+        """
         if isinstance(curr, list):
             assert isinstance(curr_pos, list)
             assert len(curr) == len(curr_pos) == 1

ultralytics/models/sam/modules/utils.py

@@ -61,7 +61,23 @@ def select_closest_cond_frames(frame_idx, cond_frame_outputs, max_cond_frame_num
 
 
 def get_1d_sine_pe(pos_inds, dim, temperature=10000):
-    """Generate 1D sinusoidal positional embeddings for given positions and dimensions."""
+    """
+    Generate 1D sinusoidal positional embeddings for given positions and dimensions.
+
+    Args:
+        pos_inds (torch.Tensor): Position indices for which to generate embeddings.
+        dim (int): Dimension of the positional embeddings. Should be an even number.
+        temperature (float): Scaling factor for the frequency of the sinusoidal functions.
+
+    Returns:
+        (torch.Tensor): Sinusoidal positional embeddings with shape (pos_inds.shape, dim).
+
+    Examples:
+        >>> pos = torch.tensor([0, 1, 2, 3])
+        >>> embeddings = get_1d_sine_pe(pos, 128)
+        >>> embeddings.shape
+        torch.Size([4, 128])
+    """
     pe_dim = dim // 2
     dim_t = torch.arange(pe_dim, dtype=torch.float32, device=pos_inds.device)
     dim_t = temperature ** (2 * (dim_t // 2) / pe_dim)
@@ -72,7 +88,30 @@ def get_1d_sine_pe(pos_inds, dim, temperature=10000):
 
 
 def init_t_xy(end_x: int, end_y: int):
-    """Initialize 1D and 2D coordinate tensors for a grid of specified dimensions."""
+    """
+    Initialize 1D and 2D coordinate tensors for a grid of specified dimensions.
+
+    This function creates coordinate tensors for a grid with dimensions end_x × end_y. It generates a linear index tensor
+    and corresponding x and y coordinate tensors.
+
+    Args:
+        end_x (int): Width of the grid (number of columns).
+        end_y (int): Height of the grid (number of rows).
+
+    Returns:
+        t (torch.Tensor): Linear indices for each position in the grid, with shape (end_x * end_y).
+        t_x (torch.Tensor): X-coordinates for each position, with shape (end_x * end_y).
+        t_y (torch.Tensor): Y-coordinates for each position, with shape (end_x * end_y).
+
+    Examples:
+        >>> t, t_x, t_y = init_t_xy(3, 2)
+        >>> print(t)
+        tensor([0., 1., 2., 3., 4., 5.])
+        >>> print(t_x)
+        tensor([0., 1., 2., 0., 1., 2.])
+        >>> print(t_y)
+        tensor([0., 0., 0., 1., 1., 1.])
+    """
     t = torch.arange(end_x * end_y, dtype=torch.float32)
     t_x = (t % end_x).float()
     t_y = torch.div(t, end_x, rounding_mode="floor").float()
@@ -80,7 +119,32 @@ def init_t_xy(end_x: int, end_y: int):
 
 
 def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 10000.0):
-    """Compute axial complex exponential positional encodings for 2D spatial positions in a grid."""
+    """
+    Compute axial complex exponential positional encodings for 2D spatial positions in a grid.
+
+    This function generates complex exponential positional encodings for a 2D grid of spatial positions,
+    using separate frequency components for the x and y dimensions.
+
+    Args:
+        dim (int): Dimension of the positional encoding.
+        end_x (int): Width of the 2D grid.
+        end_y (int): Height of the 2D grid.
+        theta (float, optional): Scaling factor for frequency computation.
+
+    Returns:
+        freqs_cis_x (torch.Tensor): Complex exponential positional encodings for x-dimension with shape
+            (end_x*end_y, dim//4).
+        freqs_cis_y (torch.Tensor): Complex exponential positional encodings for y-dimension with shape
+            (end_x*end_y, dim//4).
+
+    Examples:
+        >>> dim, end_x, end_y = 128, 8, 8
+        >>> freqs_cis_x, freqs_cis_y = compute_axial_cis(dim, end_x, end_y)
+        >>> freqs_cis_x.shape
+        torch.Size([64, 32])
+        >>> freqs_cis_y.shape
+        torch.Size([64, 32])
+    """
     freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
     freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
 
@@ -93,7 +157,22 @@ def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 10000.0):
 
 
 def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
-    """Reshape frequency tensor for broadcasting with input tensor, ensuring dimensional compatibility."""
+    """
+    Reshape frequency tensor for broadcasting with input tensor.
+
+    Reshapes a frequency tensor to ensure dimensional compatibility for broadcasting with an input tensor.
+    This function is typically used in positional encoding operations.
+
+    Args:
+        freqs_cis (torch.Tensor): Frequency tensor with shape matching the last two dimensions of x.
+        x (torch.Tensor): Input tensor to broadcast with.
+
+    Returns:
+        (torch.Tensor): Reshaped frequency tensor ready for broadcasting with the input tensor.
+
+    Raises:
+        AssertionError: If the shape of freqs_cis doesn't match the last two dimensions of x.
+    """
     ndim = x.ndim
     assert 0 <= 1 < ndim
     assert freqs_cis.shape == (x.shape[-2], x.shape[-1])
@@ -107,7 +186,31 @@ def apply_rotary_enc(
     freqs_cis: torch.Tensor,
     repeat_freqs_k: bool = False,
 ):
-    """Apply rotary positional encoding to query and key tensors using complex-valued frequency components."""
+    """
+    Apply rotary positional encoding to query and key tensors.
+
+    This function applies rotary positional encoding (RoPE) to query and key tensors using complex-valued frequency
+    components. RoPE is a technique that injects relative position information into self-attention mechanisms.
+
+    Args:
+        xq (torch.Tensor): Query tensor to encode with positional information.
+        xk (torch.Tensor): Key tensor to encode with positional information.
+        freqs_cis (torch.Tensor): Complex-valued frequency components for rotary encoding with shape matching the
+            last two dimensions of xq.
+        repeat_freqs_k (bool, optional): Whether to repeat frequency components along sequence length dimension
+            to match key sequence length.
+
+    Returns:
+        xq_out (torch.Tensor): Query tensor with rotary positional encoding applied.
+        xk_out (torch.Tensor): Key tensor with rotary positional encoding applied, or original xk if xk is empty.
+
+    Examples:
+        >>> import torch
+        >>> xq = torch.randn(2, 8, 16, 64)  # [batch, heads, seq_len, dim]
+        >>> xk = torch.randn(2, 8, 16, 64)
+        >>> freqs_cis = compute_axial_cis(64, 4, 4)  # For a 4x4 spatial grid with dim=64
+        >>> q_encoded, k_encoded = apply_rotary_enc(xq, xk, freqs_cis)
+    """
     xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
     xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) if xk.shape[-2] != 0 else None
     freqs_cis = reshape_for_broadcast(freqs_cis, xq_)

ultralytics/models/utils/loss.py

@@ -65,7 +65,25 @@ class DETRLoss(nn.Module):
         self.device = None
 
     def _get_loss_class(self, pred_scores, targets, gt_scores, num_gts, postfix=""):
-        """Compute classification loss based on predictions, target values, and ground truth scores."""
+        """
+        Compute classification loss based on predictions, target values, and ground truth scores.
+
+        Args:
+            pred_scores (torch.Tensor): Predicted class scores with shape (batch_size, num_queries, num_classes).
+            targets (torch.Tensor): Target class indices with shape (batch_size, num_queries).
+            gt_scores (torch.Tensor): Ground truth confidence scores with shape (batch_size, num_queries).
+            num_gts (int): Number of ground truth objects.
+            postfix (str, optional): String to append to the loss name for identification in multi-loss scenarios.
+
+        Returns:
+            loss_cls (torch.Tensor): Classification loss value.
+
+        Notes:
+            The function supports different classification loss types:
+            - Varifocal Loss (if self.vfl is True and num_gts > 0)
+            - Focal Loss (if self.fl is True)
+            - BCE Loss (default fallback)
+        """
         # Logits: [b, query, num_classes], gt_class: list[[n, 1]]
         name_class = f"loss_class{postfix}"
         bs, nq = pred_scores.shape[:2]
@@ -87,7 +105,25 @@ class DETRLoss(nn.Module):
         return {name_class: loss_cls.squeeze() * self.loss_gain["class"]}
 
     def _get_loss_bbox(self, pred_bboxes, gt_bboxes, postfix=""):
-        """Compute bounding box and GIoU losses for predicted and ground truth bounding boxes."""
+        """
+        Compute bounding box and GIoU losses for predicted and ground truth bounding boxes.
+
+        Args:
+            pred_bboxes (torch.Tensor): Predicted bounding boxes with shape (batch_size, num_queries, 4).
+            gt_bboxes (torch.Tensor): Ground truth bounding boxes with shape (N, 4), where N is the total
+                number of ground truth boxes.
+            postfix (str): String to append to the loss names for identification in multi-loss scenarios.
+
+        Returns:
+            loss (dict): Dictionary containing:
+                - loss_bbox{postfix} (torch.Tensor): L1 loss between predicted and ground truth boxes,
+                  scaled by the bbox loss gain.
+                - loss_giou{postfix} (torch.Tensor): GIoU loss between predicted and ground truth boxes,
+                  scaled by the giou loss gain.
+
+        Notes:
+            If no ground truth boxes are provided (empty list), zero-valued tensors are returned for both losses.
+        """
         # Boxes: [b, query, 4], gt_bbox: list[[n, 4]]
         name_bbox = f"loss_bbox{postfix}"
         name_giou = f"loss_giou{postfix}"

ultralytics/models/utils/ops.py

@@ -31,7 +31,21 @@ class HungarianMatcher(nn.Module):
     """
 
     def __init__(self, cost_gain=None, use_fl=True, with_mask=False, num_sample_points=12544, alpha=0.25, gamma=2.0):
-        """Initialize a HungarianMatcher module for optimal assignment of predicted and ground truth bounding boxes."""
+        """
+        Initialize a HungarianMatcher module for optimal assignment of predicted and ground truth bounding boxes.
+
+        The HungarianMatcher uses a cost function that considers classification scores, bounding box coordinates,
+        and optionally mask predictions to perform optimal bipartite matching between predictions and ground truths.
+
+        Args:
+            cost_gain (dict, optional): Dictionary of cost coefficients for different components of the matching cost.
+                Should contain keys 'class', 'bbox', 'giou', 'mask', and 'dice'.
+            use_fl (bool, optional): Whether to use Focal Loss for the classification cost calculation.
+            with_mask (bool, optional): Whether the model makes mask predictions.
+            num_sample_points (int, optional): Number of sample points used in mask cost calculation.
+            alpha (float, optional): Alpha factor in Focal Loss calculation.
+            gamma (float, optional): Gamma factor in Focal Loss calculation.
+        """
         super().__init__()
         if cost_gain is None:
             cost_gain = {"class": 1, "bbox": 5, "giou": 2, "mask": 1, "dice": 1}

ultralytics/models/yolo/classify/predict.py

@@ -36,7 +36,17 @@ class ClassificationPredictor(BasePredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initialize the ClassificationPredictor with the specified configuration and set task to 'classify'."""
+        """
+        Initialize the ClassificationPredictor with the specified configuration and set task to 'classify'.
+
+        This constructor initializes a ClassificationPredictor instance, which extends BasePredictor for classification
+        tasks. It ensures the task is set to 'classify' regardless of input configuration.
+
+        Args:
+            cfg (dict): Default configuration dictionary containing prediction settings. Defaults to DEFAULT_CFG.
+            overrides (dict, optional): Configuration overrides that take precedence over cfg.
+            _callbacks (list, optional): List of callback functions to be executed during prediction.
+        """
         super().__init__(cfg, overrides, _callbacks)
         self.args.task = "classify"
         self._legacy_transform_name = "ultralytics.yolo.data.augment.ToTensor"

ultralytics/models/yolo/classify/train.py

@@ -48,7 +48,23 @@ class ClassificationTrainer(BaseTrainer):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initialize a ClassificationTrainer object with optional configuration overrides and callbacks."""
+        """
+        Initialize a ClassificationTrainer object.
+
+        This constructor sets up a trainer for image classification tasks, configuring the task type and default
+        image size if not specified.
+
+        Args:
+            cfg (dict, optional): Default configuration dictionary containing training parameters.
+            overrides (dict, optional): Dictionary of parameter overrides for the default configuration.
+            _callbacks (list, optional): List of callback functions to be executed during training.
+
+        Examples:
+            >>> from ultralytics.models.yolo.classify import ClassificationTrainer
+            >>> args = dict(model="yolo11n-cls.pt", data="imagenet10", epochs=3)
+            >>> trainer = ClassificationTrainer(overrides=args)
+            >>> trainer.train()
+        """
         if overrides is None:
             overrides = {}
         overrides["task"] = "classify"

ultralytics/models/yolo/classify/val.py

@@ -49,7 +49,25 @@ class ClassificationValidator(BaseValidator):
     """
 
     def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
-        """Initialize ClassificationValidator with dataloader, save directory, and other parameters."""
+        """
+        Initialize ClassificationValidator with dataloader, save directory, and other parameters.
+
+        This validator handles the validation process for classification models, including metrics calculation,
+        confusion matrix generation, and visualization of results.
+
+        Args:
+            dataloader (torch.utils.data.DataLoader, optional): Dataloader to use for validation.
+            save_dir (str | Path, optional): Directory to save results.
+            pbar (bool, optional): Display a progress bar.
+            args (dict, optional): Arguments containing model and validation configuration.
+            _callbacks (list, optional): List of callback functions to be called during validation.
+
+        Examples:
+            >>> from ultralytics.models.yolo.classify import ClassificationValidator
+            >>> args = dict(model="yolo11n-cls.pt", data="imagenet10")
+            >>> validator = ClassificationValidator(args=args)
+            >>> validator()
+        """
         super().__init__(dataloader, save_dir, pbar, args, _callbacks)
         self.targets = None
         self.pred = None
@@ -76,13 +94,38 @@ class ClassificationValidator(BaseValidator):
         return batch
 
     def update_metrics(self, preds, batch):
-        """Update running metrics with model predictions and batch targets."""
+        """
+        Update running metrics with model predictions and batch targets.
+
+        Args:
+            preds (torch.Tensor): Model predictions, typically logits or probabilities for each class.
+            batch (dict): Batch data containing images and class labels.
+
+        This method appends the top-N predictions (sorted by confidence in descending order) to the
+        prediction list for later evaluation. N is limited to the minimum of 5 and the number of classes.
+        """
         n5 = min(len(self.names), 5)
         self.pred.append(preds.argsort(1, descending=True)[:, :n5].type(torch.int32).cpu())
         self.targets.append(batch["cls"].type(torch.int32).cpu())
 
     def finalize_metrics(self, *args, **kwargs):
-        """Finalize metrics including confusion matrix and processing speed."""
+        """
+        Finalize metrics including confusion matrix and processing speed.
+
+        This method processes the accumulated predictions and targets to generate the confusion matrix,
+        optionally plots it, and updates the metrics object with speed information.
+
+        Args:
+            *args (Any): Variable length argument list.
+            **kwargs (Any): Arbitrary keyword arguments.
+
+        Examples:
+            >>> validator = ClassificationValidator()
+            >>> validator.pred = [torch.tensor([[0, 1, 2]])]  # Top-3 predictions for one sample
+            >>> validator.targets = [torch.tensor([0])]  # Ground truth class
+            >>> validator.finalize_metrics()
+            >>> print(validator.metrics.confusion_matrix)  # Access the confusion matrix
+        """
         self.confusion_matrix.process_cls_preds(self.pred, self.targets)
         if self.args.plots:
             for normalize in True, False:
@@ -107,7 +150,16 @@ class ClassificationValidator(BaseValidator):
         return ClassificationDataset(root=img_path, args=self.args, augment=False, prefix=self.args.split)
 
     def get_dataloader(self, dataset_path, batch_size):
-        """Build and return a data loader for classification validation."""
+        """
+        Build and return a data loader for classification validation.
+
+        Args:
+            dataset_path (str | Path): Path to the dataset directory.
+            batch_size (int): Number of samples per batch.
+
+        Returns:
+            (torch.utils.data.DataLoader): DataLoader object for the classification validation dataset.
+        """
         dataset = self.build_dataset(dataset_path)
         return build_dataloader(dataset, batch_size, self.args.workers, rank=-1)
 
@@ -117,7 +169,18 @@ class ClassificationValidator(BaseValidator):
         LOGGER.info(pf % ("all", self.metrics.top1, self.metrics.top5))
 
     def plot_val_samples(self, batch, ni):
-        """Plot validation image samples with their ground truth labels."""
+        """
+        Plot validation image samples with their ground truth labels.
+
+        Args:
+            batch (dict): Dictionary containing batch data with 'img' (images) and 'cls' (class labels).
+            ni (int): Batch index used for naming the output file.
+
+        Examples:
+            >>> validator = ClassificationValidator()
+            >>> batch = {"img": torch.rand(16, 3, 224, 224), "cls": torch.randint(0, 10, (16,))}
+            >>> validator.plot_val_samples(batch, 0)
+        """
         plot_images(
             images=batch["img"],
             batch_idx=torch.arange(len(batch["img"])),
@@ -128,7 +191,20 @@ class ClassificationValidator(BaseValidator):
         )
 
     def plot_predictions(self, batch, preds, ni):
-        """Plot images with their predicted class labels and save the visualization."""
+        """
+        Plot images with their predicted class labels and save the visualization.
+
+        Args:
+            batch (dict): Batch data containing images and other information.
+            preds (torch.Tensor): Model predictions with shape (batch_size, num_classes).
+            ni (int): Batch index used for naming the output file.
+
+        Examples:
+            >>> validator = ClassificationValidator()
+            >>> batch = {"img": torch.rand(16, 3, 224, 224)}
+            >>> preds = torch.rand(16, 10)  # 16 images, 10 classes
+            >>> validator.plot_predictions(batch, preds, 0)
+        """
         plot_images(
             batch["img"],
             batch_idx=torch.arange(len(batch["img"])),

ultralytics/models/yolo/detect/predict.py

@@ -31,7 +31,26 @@ class DetectionPredictor(BasePredictor):
     """
 
     def postprocess(self, preds, img, orig_imgs, **kwargs):
-        """Post-processes predictions and returns a list of Results objects."""
+        """
+        Post-process predictions and return a list of Results objects.
+
+        This method applies non-maximum suppression to raw model predictions and prepares them for visualization and
+        further analysis.
+
+        Args:
+            preds (torch.Tensor): Raw predictions from the model.
+            img (torch.Tensor): Processed input image tensor in model input format.
+            orig_imgs (torch.Tensor | list): Original input images before preprocessing.
+            **kwargs (Any): Additional keyword arguments.
+
+        Returns:
+            (list): List of Results objects containing the post-processed predictions.
+
+        Examples:
+            >>> predictor = DetectionPredictor(overrides=dict(model="yolov8n.pt"))
+            >>> results = predictor.predict("path/to/image.jpg")
+            >>> processed_results = predictor.postprocess(preds, img, orig_imgs)
+        """
         preds = ops.non_max_suppression(
             preds,
             self.args.conf,

ultralytics/models/yolo/model.py

@@ -2,6 +2,7 @@
 
 from pathlib import Path
 
+from ultralytics.data.build import load_inference_source
 from ultralytics.engine.model import Model
 from ultralytics.models import yolo
 from ultralytics.nn.tasks import (
@@ -21,7 +22,24 @@ class YOLO(Model):
     """YOLO (You Only Look Once) object detection model."""
 
     def __init__(self, model="yolo11n.pt", task=None, verbose=False):
-        """Initialize YOLO model, switching to YOLOWorld/YOLOE if model filename contains '-world'/'yoloe'."""
+        """
+        Initialize a YOLO model.
+
+        This constructor initializes a YOLO model, automatically switching to specialized model types
+        (YOLOWorld or YOLOE) based on the model filename.
+
+        Args:
+            model (str | Path): Model name or path to model file, i.e. 'yolo11n.pt', 'yolov8n.yaml'.
+            task (str | None): YOLO task specification, i.e. 'detect', 'segment', 'classify', 'pose', 'obb'.
+                Defaults to auto-detection based on model.
+            verbose (bool): Display model info on load.
+
+        Examples:
+            >>> from ultralytics import YOLO
+            >>> model = YOLO("yolov8n.pt")  # load a pretrained YOLOv8n detection model
+            >>> model = YOLO("yolov8n-seg.pt")  # load a pretrained YOLOv8n segmentation model
+            >>> model = YOLO("yolo11n.pt")  # load a pretrained YOLOv11n detection model
+        """
         path = Path(model)
         if "-world" in path.stem and path.suffix in {".pt", ".yaml", ".yml"}:  # if YOLOWorld PyTorch model
             new_instance = YOLOWorld(path, verbose=verbose)
@@ -165,12 +183,46 @@ class YOLOE(Model):
         return self.model.get_text_pe(texts)
 
     def get_visual_pe(self, img, visual):
-        """Get visual positional embeddings for the given image and visual features."""
+        """
+        Get visual positional embeddings for the given image and visual features.
+
+        This method extracts positional embeddings from visual features based on the input image. It requires
+        that the model is an instance of YOLOEModel.
+
+        Args:
+            img (torch.Tensor): Input image tensor.
+            visual (torch.Tensor): Visual features extracted from the image.
+
+        Returns:
+            (torch.Tensor): Visual positional embeddings.
+
+        Examples:
+            >>> model = YOLOE("yoloe-v8s.pt")
+            >>> img = torch.rand(1, 3, 640, 640)
+            >>> visual_features = model.model.backbone(img)
+            >>> pe = model.get_visual_pe(img, visual_features)
+        """
         assert isinstance(self.model, YOLOEModel)
         return self.model.get_visual_pe(img, visual)
 
     def set_vocab(self, vocab, names):
-        """Set vocabulary and class names for the model."""
+        """
+        Set vocabulary and class names for the YOLOE model.
+
+        This method configures the vocabulary and class names used by the model for text processing and
+        classification tasks. The model must be an instance of YOLOEModel.
+
+        Args:
+            vocab (list): Vocabulary list containing tokens or words used by the model for text processing.
+            names (list): List of class names that the model can detect or classify.
+
+        Raises:
+            AssertionError: If the model is not an instance of YOLOEModel.
+
+        Examples:
+            >>> model = YOLOE("yoloe-v8s.pt")
+            >>> model.set_vocab(["person", "car", "dog"], ["person", "car", "dog"])
+        """
         assert isinstance(self.model, YOLOEModel)
         self.model.set_vocab(vocab, names=names)
 
@@ -267,7 +319,14 @@ class YOLOE(Model):
                 f"{len(visual_prompts['cls'])} respectively"
             )
         self.predictor = (predictor or self._smart_load("predictor"))(
-            overrides={"task": "segment", "mode": "predict", "save": False, "verbose": False}, _callbacks=self.callbacks
+            overrides={
+                "task": self.model.task,
+                "mode": "predict",
+                "save": False,
+                "verbose": refer_image is None,
+                "batch": 1,
+            },
+            _callbacks=self.callbacks,
         )
 
         if len(visual_prompts):
@@ -281,6 +340,12 @@ class YOLOE(Model):
             self.predictor.set_prompts(visual_prompts.copy())
 
         self.predictor.setup_model(model=self.model)
+
+        if refer_image is None and source:
+            dataset = load_inference_source(source)
+            if dataset.mode in {"video", "stream"}:
+                # NOTE: set the first frame as refer image for videos/streams inference
+                refer_image = next(iter(dataset))[1][0]
         if refer_image is not None and len(visual_prompts):
             vpe = self.predictor.get_vpe(refer_image)
             self.model.set_classes(self.model.names, vpe)

ultralytics/models/yolo/obb/predict.py

@@ -27,7 +27,22 @@ class OBBPredictor(DetectionPredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initialize OBBPredictor with optional model and data configuration overrides."""
+        """
+        Initialize OBBPredictor with optional model and data configuration overrides.
+
+        This constructor sets up an OBBPredictor instance for oriented bounding box detection tasks.
+
+        Args:
+            cfg (dict, optional): Default configuration for the predictor.
+            overrides (dict, optional): Configuration overrides that take precedence over the default config.
+            _callbacks (list, optional): List of callback functions to be invoked during prediction.
+
+        Examples:
+            >>> from ultralytics.utils import ASSETS
+            >>> from ultralytics.models.yolo.obb import OBBPredictor
+            >>> args = dict(model="yolo11n-obb.pt", source=ASSETS)
+            >>> predictor = OBBPredictor(overrides=args)
+        """
         super().__init__(cfg, overrides, _callbacks)
         self.args.task = "obb"