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

ultralytics/models/yolo/segment/val.py

@@ -18,6 +18,16 @@ class SegmentationValidator(DetectionValidator):
     """
     A class extending the DetectionValidator class for validation based on a segmentation model.
 
+    This validator handles the evaluation of segmentation models, processing both bounding box and mask predictions
+    to compute metrics such as mAP for both detection and segmentation tasks.
+
+    Attributes:
+        plot_masks (List): List to store masks for plotting.
+        process (callable): Function to process masks based on save_json and save_txt flags.
+        args (namespace): Arguments for the validator.
+        metrics (SegmentMetrics): Metrics calculator for segmentation tasks.
+        stats (Dict): Dictionary to store statistics during validation.
+
     Examples:
         >>> from ultralytics.models.yolo.segment import SegmentationValidator
         >>> args = dict(model="yolo11n-seg.pt", data="coco8-seg.yaml")
@@ -26,7 +36,16 @@ class SegmentationValidator(DetectionValidator):
     """
 
     def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
-        """Initialize SegmentationValidator and set task to 'segment', metrics to SegmentMetrics."""
+        """
+        Initialize SegmentationValidator and set task to 'segment', metrics to SegmentMetrics.
+
+        Args:
+            dataloader (torch.utils.data.DataLoader, optional): Dataloader to use for validation.
+            save_dir (Path, optional): Directory to save results.
+            pbar (Any, optional): Progress bar for displaying progress.
+            args (namespace, optional): Arguments for the validator.
+            _callbacks (List, optional): List of callback functions.
+        """
         super().__init__(dataloader, save_dir, pbar, args, _callbacks)
         self.plot_masks = None
         self.process = None
@@ -34,13 +53,18 @@ class SegmentationValidator(DetectionValidator):
         self.metrics = SegmentMetrics(save_dir=self.save_dir)
 
     def preprocess(self, batch):
-        """Preprocesses batch by converting masks to float and sending to device."""
+        """Preprocess batch by converting masks to float and sending to device."""
         batch = super().preprocess(batch)
         batch["masks"] = batch["masks"].to(self.device).float()
         return batch
 
     def init_metrics(self, model):
-        """Initialize metrics and select mask processing function based on save_json flag."""
+        """
+        Initialize metrics and select mask processing function based on save_json flag.
+
+        Args:
+            model (torch.nn.Module): Model to validate.
+        """
         super().init_metrics(model)
         self.plot_masks = []
         if self.args.save_json:
@@ -66,26 +90,61 @@ class SegmentationValidator(DetectionValidator):
         )
 
     def postprocess(self, preds):
-        """Post-processes YOLO predictions and returns output detections with proto."""
+        """
+        Post-process YOLO predictions and return output detections with proto.
+
+        Args:
+            preds (List): Raw predictions from the model.
+
+        Returns:
+            p (torch.Tensor): Processed detection predictions.
+            proto (torch.Tensor): Prototype masks for segmentation.
+        """
         p = super().postprocess(preds[0])
         proto = preds[1][-1] if len(preds[1]) == 3 else preds[1]  # second output is len 3 if pt, but only 1 if exported
         return p, proto
 
     def _prepare_batch(self, si, batch):
-        """Prepares a batch for training or inference by processing images and targets."""
+        """
+        Prepare a batch for training or inference by processing images and targets.
+
+        Args:
+            si (int): Batch index.
+            batch (Dict): Batch data containing images and targets.
+
+        Returns:
+            (Dict): Prepared batch with processed images and targets.
+        """
         prepared_batch = super()._prepare_batch(si, batch)
         midx = [si] if self.args.overlap_mask else batch["batch_idx"] == si
         prepared_batch["masks"] = batch["masks"][midx]
         return prepared_batch
 
     def _prepare_pred(self, pred, pbatch, proto):
-        """Prepares a batch for training or inference by processing images and targets."""
+        """
+        Prepare predictions for evaluation by processing bounding boxes and masks.
+
+        Args:
+            pred (torch.Tensor): Raw predictions from the model.
+            pbatch (Dict): Prepared batch data.
+            proto (torch.Tensor): Prototype masks for segmentation.
+
+        Returns:
+            predn (torch.Tensor): Processed bounding box predictions.
+            pred_masks (torch.Tensor): Processed mask predictions.
+        """
         predn = super()._prepare_pred(pred, pbatch)
         pred_masks = self.process(proto, pred[:, 6:], pred[:, :4], shape=pbatch["imgsz"])
         return predn, pred_masks
 
     def update_metrics(self, preds, batch):
-        """Metrics."""
+        """
+        Update metrics with the current batch predictions and targets.
+
+        Args:
+            preds (List): Predictions from the model.
+            batch (Dict): Batch data containing images and targets.
+        """
         for si, (pred, proto) in enumerate(zip(preds[0], preds[1])):
             self.seen += 1
             npr = len(pred)
@@ -154,7 +213,7 @@ class SegmentationValidator(DetectionValidator):
                 )
 
     def finalize_metrics(self, *args, **kwargs):
-        """Sets speed and confusion matrix for evaluation metrics."""
+        """Set speed and confusion matrix for evaluation metrics."""
         self.metrics.speed = self.speed
         self.metrics.confusion_matrix = self.confusion_matrix
 
@@ -168,9 +227,9 @@ class SegmentationValidator(DetectionValidator):
             gt_bboxes (torch.Tensor): Tensor of shape (M, 4) representing ground truth bounding box coordinates.
                 Each row is of the format [x1, y1, x2, y2].
             gt_cls (torch.Tensor): Tensor of shape (M,) representing ground truth class indices.
-            pred_masks (torch.Tensor | None): Tensor representing predicted masks, if available. The shape should
+            pred_masks (torch.Tensor, optional): Tensor representing predicted masks, if available. The shape should
                 match the ground truth masks.
-            gt_masks (torch.Tensor | None): Tensor of shape (M, H, W) representing ground truth masks, if available.
+            gt_masks (torch.Tensor, optional): Tensor of shape (M, H, W) representing ground truth masks, if available.
             overlap (bool): Flag indicating if overlapping masks should be considered.
             masks (bool): Flag indicating if the batch contains mask data.
 
@@ -203,7 +262,13 @@ class SegmentationValidator(DetectionValidator):
         return self.match_predictions(detections[:, 5], gt_cls, iou)
 
     def plot_val_samples(self, batch, ni):
-        """Plots validation samples with bounding box labels."""
+        """
+        Plot validation samples with bounding box labels and masks.
+
+        Args:
+            batch (Dict): Batch data containing images and targets.
+            ni (int): Batch index.
+        """
         plot_images(
             batch["img"],
             batch["batch_idx"],
@@ -217,7 +282,14 @@ class SegmentationValidator(DetectionValidator):
         )
 
     def plot_predictions(self, batch, preds, ni):
-        """Plots batch predictions with masks and bounding boxes."""
+        """
+        Plot batch predictions with masks and bounding boxes.
+
+        Args:
+            batch (Dict): Batch data containing images.
+            preds (List): Predictions from the model.
+            ni (int): Batch index.
+        """
         plot_images(
             batch["img"],
             *output_to_target(preds[0], max_det=15),  # not set to self.args.max_det due to slow plotting speed
@@ -230,7 +302,16 @@ class SegmentationValidator(DetectionValidator):
         self.plot_masks.clear()
 
     def save_one_txt(self, predn, pred_masks, save_conf, shape, file):
-        """Save YOLO detections to a txt file in normalized coordinates in a specific format."""
+        """
+        Save YOLO detections to a txt file in normalized coordinates in a specific format.
+
+        Args:
+            predn (torch.Tensor): Predictions in the format [x1, y1, x2, y2, conf, cls].
+            pred_masks (torch.Tensor): Predicted masks.
+            save_conf (bool): Whether to save confidence scores.
+            shape (Tuple): Original image shape.
+            file (Path): File path to save the detections.
+        """
         from ultralytics.engine.results import Results
 
         Results(
@@ -243,7 +324,12 @@ class SegmentationValidator(DetectionValidator):
 
     def pred_to_json(self, predn, filename, pred_masks):
         """
-        Save one JSON result.
+        Save one JSON result for COCO evaluation.
+
+        Args:
+            predn (torch.Tensor): Predictions in the format [x1, y1, x2, y2, conf, cls].
+            filename (str): Image filename.
+            pred_masks (numpy.ndarray): Predicted masks.
 
         Examples:
              >>> result = {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}

ultralytics/models/yolo/world/train.py

@@ -10,9 +10,9 @@ from ultralytics.utils.torch_utils import de_parallel
 
 
 def on_pretrain_routine_end(trainer):
-    """Callback."""
+    """Callback to set up model classes and text encoder at the end of the pretrain routine."""
     if RANK in {-1, 0}:
-        # NOTE: for evaluation
+        # Set class names for evaluation
         names = [name.split("/")[0] for name in list(trainer.test_loader.dataset.data["names"].values())]
         de_parallel(trainer.ema.ema).set_classes(names, cache_clip_model=False)
     device = next(trainer.model.parameters()).device
@@ -25,6 +25,16 @@ class WorldTrainer(yolo.detect.DetectionTrainer):
     """
     A class to fine-tune a world model on a close-set dataset.
 
+    This trainer extends the DetectionTrainer to support training YOLO World models, which combine
+    visual and textual features for improved object detection and understanding.
+
+    Attributes:
+        clip (module): The CLIP module for text-image understanding.
+        text_model (module): The text encoder model from CLIP.
+        model (WorldModel): The YOLO World model being trained.
+        data (Dict): Dataset configuration containing class information.
+        args (Dict): Training arguments and configuration.
+
     Examples:
         >>> from ultralytics.models.yolo.world import WorldModel
         >>> args = dict(model="yolov8s-world.pt", data="coco8.yaml", epochs=3)
@@ -33,7 +43,14 @@ class WorldTrainer(yolo.detect.DetectionTrainer):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initialize a WorldTrainer object with given arguments."""
+        """
+        Initialize a WorldTrainer object with given arguments.
+
+        Args:
+            cfg (Dict): Configuration for the trainer.
+            overrides (Dict, optional): Configuration overrides.
+            _callbacks (List, optional): List of callback functions.
+        """
         if overrides is None:
             overrides = {}
         super().__init__(cfg, overrides, _callbacks)
@@ -47,7 +64,17 @@ class WorldTrainer(yolo.detect.DetectionTrainer):
         self.clip = clip
 
     def get_model(self, cfg=None, weights=None, verbose=True):
-        """Return WorldModel initialized with specified config and weights."""
+        """
+        Return WorldModel initialized with specified config and weights.
+
+        Args:
+            cfg (Dict | str, optional): Model configuration.
+            weights (str, optional): Path to pretrained weights.
+            verbose (bool): Whether to display model info.
+
+        Returns:
+            (WorldModel): Initialized WorldModel.
+        """
         # NOTE: This `nc` here is the max number of different text samples in one image, rather than the actual `nc`.
         # NOTE: Following the official config, nc hard-coded to 80 for now.
         model = WorldModel(
@@ -64,12 +91,15 @@ class WorldTrainer(yolo.detect.DetectionTrainer):
 
     def build_dataset(self, img_path, mode="train", batch=None):
         """
-        Build YOLO Dataset.
+        Build YOLO Dataset for training or validation.
 
         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:
+            (Dataset): YOLO dataset configured for training or validation.
         """
         gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
         return build_yolo_dataset(
@@ -77,10 +107,10 @@ class WorldTrainer(yolo.detect.DetectionTrainer):
         )
 
     def preprocess_batch(self, batch):
-        """Preprocesses a batch of images for YOLOWorld training, adjusting formatting and dimensions as needed."""
+        """Preprocess a batch of images and text for YOLOWorld training."""
         batch = super().preprocess_batch(batch)
 
-        # NOTE: add text features
+        # Add text features
         texts = list(itertools.chain(*batch["texts"]))
         text_token = self.clip.tokenize(texts).to(batch["img"].device)
         txt_feats = self.text_model.encode_text(text_token).to(dtype=batch["img"].dtype)  # torch.float32

ultralytics/models/yolo/world/train_world.py

@@ -9,7 +9,15 @@ from ultralytics.utils.torch_utils import de_parallel
 
 class WorldTrainerFromScratch(WorldTrainer):
     """
-    A class extending the WorldTrainer class for training a world model from scratch on open-set dataset.
+    A class extending the WorldTrainer for training a world model from scratch on open-set datasets.
+
+    This trainer specializes in handling mixed datasets including both object detection and grounding datasets,
+    supporting training YOLO-World models with combined vision-language capabilities.
+
+    Attributes:
+        cfg (Dict): Configuration dictionary with default parameters for model training.
+        overrides (Dict): Dictionary of parameter overrides to customize the configuration.
+        _callbacks (List): List of callback functions to be executed during different stages of training.
 
     Examples:
         >>> from ultralytics.models.yolo.world.train_world import WorldTrainerFromScratch
@@ -35,19 +43,25 @@ class WorldTrainerFromScratch(WorldTrainer):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initialize a WorldTrainer object with given arguments."""
+        """Initialize a WorldTrainerFromScratch object with given configuration and callbacks."""
         if overrides is None:
             overrides = {}
         super().__init__(cfg, overrides, _callbacks)
 
     def build_dataset(self, img_path, mode="train", batch=None):
         """
-        Build YOLO Dataset.
+        Build YOLO Dataset for training or validation.
+
+        This method constructs appropriate datasets based on the mode and input paths, handling both
+        standard YOLO datasets and grounding datasets with different formats.
 
         Args:
-            img_path (List[str] | 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.
+            img_path (List[str] | str): Path to the folder containing images or list of paths.
+            mode (str): 'train' mode or 'val' mode, allowing customized augmentations for each mode.
+            batch (int, optional): Size of batches, used for rectangular training/validation.
+
+        Returns:
+            (YOLOConcatDataset | Dataset): The constructed dataset for training or validation.
         """
         gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
         if mode != "train":
@@ -62,9 +76,17 @@ class WorldTrainerFromScratch(WorldTrainer):
 
     def get_dataset(self):
         """
-        Get train, val path from data dict if it exists.
+        Get train and validation paths from data dictionary.
+
+        Processes the data configuration to extract paths for training and validation datasets,
+        handling both YOLO detection datasets and grounding datasets.
 
-        Returns None if data format is not recognized.
+        Returns:
+            (str): Train dataset path.
+            (str): Validation dataset path.
+
+        Raises:
+            AssertionError: If train or validation datasets are not found, or if validation has multiple datasets.
         """
         final_data = {}
         data_yaml = self.args.data
@@ -94,11 +116,18 @@ class WorldTrainerFromScratch(WorldTrainer):
         return final_data["train"], final_data["val"][0]
 
     def plot_training_labels(self):
-        """DO NOT plot labels."""
+        """Do not plot labels for YOLO-World training."""
         pass
 
     def final_eval(self):
-        """Performs final evaluation and validation for object detection YOLO-World model."""
+        """
+        Perform final evaluation and validation for the YOLO-World model.
+
+        Configures the validator with appropriate dataset and split information before running evaluation.
+
+        Returns:
+            (Dict): Dictionary containing evaluation metrics and results.
+        """
         val = self.args.data["val"]["yolo_data"][0]
         self.validator.args.data = val
         self.validator.args.split = "minival" if isinstance(val, str) and "lvis" in val else "val"

ultralytics/nn/autobackend.py

@@ -19,11 +19,7 @@ from ultralytics.utils.downloads import attempt_download_asset, is_url
 
 
 def check_class_names(names):
-    """
-    Check class names.
-
-    Map imagenet class codes to human-readable names if required. Convert lists to dicts.
-    """
+    """Check class names and convert to dict format if needed."""
     if isinstance(names, list):  # names is a list
         names = dict(enumerate(names))  # convert to dict
     if isinstance(names, dict):
@@ -78,8 +74,23 @@ class AutoBackend(nn.Module):
             | IMX                   | *_imx_model/      |
             | RKNN                  | *_rknn_model/     |
 
-    This class offers dynamic backend switching capabilities based on the input model format, making it easier to deploy
-    models across various platforms.
+    Attributes:
+        model (torch.nn.Module): The loaded YOLO model.
+        device (torch.device): The device (CPU or GPU) on which the model is loaded.
+        task (str): The type of task the model performs (detect, segment, classify, pose).
+        names (Dict): A dictionary of class names that the model can detect.
+        stride (int): The model stride, typically 32 for YOLO models.
+        fp16 (bool): Whether the model uses half-precision (FP16) inference.
+
+    Methods:
+        forward: Run inference on an input image.
+        from_numpy: Convert numpy array to tensor.
+        warmup: Warm up the model with a dummy input.
+        _model_type: Determine the model type from file path.
+
+    Examples:
+        >>> model = AutoBackend(weights="yolov8n.pt", device="cuda")
+        >>> results = model(img)
     """
 
     @torch.no_grad()
@@ -101,7 +112,7 @@ class AutoBackend(nn.Module):
             weights (str | torch.nn.Module): Path to the model weights file or a module instance. Defaults to 'yolo11n.pt'.
             device (torch.device): Device to run the model on. Defaults to CPU.
             dnn (bool): Use OpenCV DNN module for ONNX inference. Defaults to False.
-            data (str | Path | optional): Path to the additional data.yaml file containing class names. Optional.
+            data (str | Path | optional): Path to the additional data.yaml file containing class names.
             fp16 (bool): Enable half-precision inference. Supported only on specific backends. Defaults to False.
             batch (int): Batch-size to assume for inference.
             fuse (bool): Fuse Conv2D + BatchNorm layers for optimization. Defaults to True.
@@ -539,12 +550,12 @@ class AutoBackend(nn.Module):
 
         Args:
             im (torch.Tensor): The image tensor to perform inference on.
-            augment (bool): whether to perform data augmentation during inference, defaults to False
-            visualize (bool): whether to visualize the output predictions, defaults to False
-            embed (list, optional): A list of feature vectors/embeddings to return.
+            augment (bool): Whether to perform data augmentation during inference. Defaults to False.
+            visualize (bool): Whether to visualize the output predictions. Defaults to False.
+            embed (List, optional): A list of feature vectors/embeddings to return.
 
         Returns:
-            (tuple): Tuple containing the raw output tensor, and processed output for visualization (if visualize=True)
+            (torch.Tensor | List[torch.Tensor]): The raw output tensor(s) from the model.
         """
         b, ch, h, w = im.shape  # batch, channel, height, width
         if self.fp16 and im.dtype != torch.float16:
@@ -776,10 +787,13 @@ class AutoBackend(nn.Module):
     def _model_type(p="path/to/model.pt"):
         """
         Takes a path to a model file and returns the model type. Possibles types are pt, jit, onnx, xml, engine, coreml,
-        saved_model, pb, tflite, edgetpu, tfjs, ncnn or paddle.
+        saved_model, pb, tflite, edgetpu, tfjs, ncnn, mnn, imx or paddle.
 
         Args:
-            p (str): path to the model file. Defaults to path/to/model.pt
+            p (str): Path to the model file. Defaults to path/to/model.pt
+
+        Returns:
+            (List[bool]): List of booleans indicating the model type.
 
         Examples:
             >>> model = AutoBackend(weights="path/to/model.onnx")

ultralytics/nn/modules/__init__.py

@@ -2,6 +2,9 @@
 """
 Ultralytics modules.
 
+This module provides access to various neural network components used in Ultralytics models, including convolution blocks,
+attention mechanisms, transformer components, and detection/segmentation heads.
+
 Examples:
     Visualize a module with Netron.
     >>> from ultralytics.nn.modules import *

ultralytics/nn/modules/activation.py

@@ -6,10 +6,19 @@ import torch.nn as nn
 
 
 class AGLU(nn.Module):
-    """Unified activation function module from https://github.com/kostas1515/AGLU."""
+    """
+    Unified activation function module from https://github.com/kostas1515/AGLU.
+
+    This class implements a parameterized activation function with learnable parameters lambda and kappa.
+
+    Attributes:
+        act (nn.Softplus): Softplus activation function with negative beta.
+        lambd (nn.Parameter): Learnable lambda parameter initialized with uniform distribution.
+        kappa (nn.Parameter): Learnable kappa parameter initialized with uniform distribution.
+    """
 
     def __init__(self, device=None, dtype=None) -> None:
-        """Initialize the Unified activation function."""
+        """Initialize the Unified activation function with learnable parameters."""
         super().__init__()
         self.act = nn.Softplus(beta=-1.0)
         self.lambd = nn.Parameter(nn.init.uniform_(torch.empty(1, device=device, dtype=dtype)))  # lambda parameter
@@ -17,5 +26,5 @@ class AGLU(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Compute the forward pass of the Unified activation function."""
-        lam = torch.clamp(self.lambd, min=0.0001)
+        lam = torch.clamp(self.lambd, min=0.0001)  # Clamp lambda to avoid division by zero
         return torch.exp((1 / lam) * self.act((self.kappa * x) - torch.log(lam)))