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

ultralytics/models/yolo/model.py

@@ -93,7 +93,7 @@ class YOLOWorld(Model):
 
     def set_classes(self, classes):
         """
-        Set classes.
+        Set the model's class names for detection.
 
         Args:
             classes (List(str)): A list of categories i.e. ["person"].
@@ -106,6 +106,5 @@ class YOLOWorld(Model):
         self.model.names = classes
 
         # Reset method class names
-        # self.predictor = None  # reset predictor otherwise old names remain
         if self.predictor:
             self.predictor.model.names = classes

ultralytics/models/yolo/obb/predict.py

@@ -11,6 +11,13 @@ class OBBPredictor(DetectionPredictor):
     """
     A class extending the DetectionPredictor class for prediction based on an Oriented Bounding Box (OBB) model.
 
+    This predictor handles oriented bounding box detection tasks, processing images and returning results with rotated
+    bounding boxes.
+
+    Attributes:
+        args (namespace): Configuration arguments for the predictor.
+        model (torch.nn.Module): The loaded YOLO OBB model.
+
     Examples:
         >>> from ultralytics.utils import ASSETS
         >>> from ultralytics.models.yolo.obb import OBBPredictor
@@ -20,17 +27,18 @@ class OBBPredictor(DetectionPredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initializes OBBPredictor with optional model and data configuration overrides."""
+        """Initialize OBBPredictor with optional model and data configuration overrides."""
         super().__init__(cfg, overrides, _callbacks)
         self.args.task = "obb"
 
     def construct_result(self, pred, img, orig_img, img_path):
         """
-        Constructs the result object from the prediction.
+        Construct the result object from the prediction.
 
         Args:
-            pred (torch.Tensor): The predicted bounding boxes, scores, and rotation angles.
-            img (torch.Tensor): The image after preprocessing.
+            pred (torch.Tensor): The predicted bounding boxes, scores, and rotation angles with shape (N, 6) where
+                the last dimension contains [x, y, w, h, confidence, class_id, angle].
+            img (torch.Tensor): The image after preprocessing with shape (B, C, H, W).
             orig_img (np.ndarray): The original image before preprocessing.
             img_path (str): The path to the original image.
 

ultralytics/models/yolo/obb/train.py

@@ -11,6 +11,13 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
     """
     A class extending the DetectionTrainer class for training based on an Oriented Bounding Box (OBB) model.
 
+    Attributes:
+        loss_names (Tuple[str]): Names of the loss components used during training.
+
+    Methods:
+        get_model: Return OBBModel initialized with specified config and weights.
+        get_validator: Return an instance of OBBValidator for validation of YOLO model.
+
     Examples:
         >>> from ultralytics.models.yolo.obb import OBBTrainer
         >>> args = dict(model="yolo11n-obb.pt", data="dota8.yaml", epochs=3)

ultralytics/models/yolo/obb/val.py

@@ -14,6 +14,24 @@ class OBBValidator(DetectionValidator):
     """
     A class extending the DetectionValidator class for validation based on an Oriented Bounding Box (OBB) model.
 
+    This validator specializes in evaluating models that predict rotated bounding boxes, commonly used for aerial and
+    satellite imagery where objects can appear at various orientations.
+
+    Attributes:
+        args (Dict): Configuration arguments for the validator.
+        metrics (OBBMetrics): Metrics object for evaluating OBB model performance.
+        is_dota (bool): Flag indicating whether the validation dataset is in DOTA format.
+
+    Methods:
+        init_metrics: Initialize evaluation metrics for YOLO.
+        _process_batch: Process batch of detections and ground truth boxes to compute IoU matrix.
+        _prepare_batch: Prepare batch data for OBB validation.
+        _prepare_pred: Prepare predictions with scaled and padded bounding boxes.
+        plot_predictions: Plot predicted bounding boxes on input images.
+        pred_to_json: Serialize YOLO predictions to COCO json format.
+        save_one_txt: Save YOLO detections to a txt file in normalized coordinates.
+        eval_json: Evaluate YOLO output in JSON format and return performance statistics.
+
     Examples:
         >>> from ultralytics.models.yolo.obb import OBBValidator
         >>> args = dict(model="yolo11n-obb.pt", data="dota8.yaml")
@@ -31,7 +49,7 @@ class OBBValidator(DetectionValidator):
         """Initialize evaluation metrics for YOLO."""
         super().init_metrics(model)
         val = self.data.get(self.args.split, "")  # validation path
-        self.is_dota = isinstance(val, str) and "DOTA" in val  # is COCO
+        self.is_dota = isinstance(val, str) and "DOTA" in val  # check if dataset is DOTA format
 
     def _process_batch(self, detections, gt_bboxes, gt_cls):
         """
@@ -61,7 +79,7 @@ class OBBValidator(DetectionValidator):
         return self.match_predictions(detections[:, 5], gt_cls, iou)
 
     def _prepare_batch(self, si, batch):
-        """Prepares and returns a batch for OBB validation."""
+        """Prepare batch data for OBB validation with proper scaling and formatting."""
         idx = batch["batch_idx"] == si
         cls = batch["cls"][idx].squeeze(-1)
         bbox = batch["bboxes"][idx]
@@ -74,7 +92,7 @@ class OBBValidator(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 for OBB validation with scaled and padded bounding boxes."""
+        """Prepare predictions by scaling bounding boxes to original image dimensions."""
         predn = pred.clone()
         ops.scale_boxes(
             pbatch["imgsz"], predn[:, :4], pbatch["ori_shape"], ratio_pad=pbatch["ratio_pad"], xywh=True
@@ -82,7 +100,7 @@ class OBBValidator(DetectionValidator):
         return predn
 
     def plot_predictions(self, batch, preds, ni):
-        """Plots predicted bounding boxes on input images and saves the result."""
+        """Plot predicted bounding boxes on input images and save the result."""
         plot_images(
             batch["img"],
             *output_to_rotated_target(preds, max_det=self.args.max_det),
@@ -93,7 +111,7 @@ class OBBValidator(DetectionValidator):
         )  # pred
 
     def pred_to_json(self, predn, filename):
-        """Serialize YOLO predictions to COCO json format."""
+        """Convert YOLO predictions to COCO JSON format with rotated bounding box information."""
         stem = Path(filename).stem
         image_id = int(stem) if stem.isnumeric() else stem
         rbox = torch.cat([predn[:, :4], predn[:, -1:]], dim=-1)
@@ -110,7 +128,7 @@ class OBBValidator(DetectionValidator):
             )
 
     def save_one_txt(self, predn, 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 using the Results class."""
         import numpy as np
 
         from ultralytics.engine.results import Results
@@ -126,7 +144,7 @@ class OBBValidator(DetectionValidator):
         ).save_txt(file, save_conf=save_conf)
 
     def eval_json(self, stats):
-        """Evaluates YOLO output in JSON format and returns performance statistics."""
+        """Evaluate YOLO output in JSON format and save predictions in DOTA format."""
         if self.args.save_json and self.is_dota and len(self.jdict):
             import json
             import re

ultralytics/models/yolo/pose/predict.py

@@ -8,6 +8,16 @@ class PosePredictor(DetectionPredictor):
     """
     A class extending the DetectionPredictor class for prediction based on a pose model.
 
+    This class specializes in pose estimation, handling keypoints detection alongside standard object detection
+    capabilities inherited from DetectionPredictor.
+
+    Attributes:
+        args (namespace): Configuration arguments for the predictor.
+        model (torch.nn.Module): The loaded YOLO pose model with keypoint detection capabilities.
+
+    Methods:
+        construct_result: Constructs the result object from the prediction, including keypoints.
+
     Examples:
         >>> from ultralytics.utils import ASSETS
         >>> from ultralytics.models.yolo.pose import PosePredictor
@@ -17,7 +27,7 @@ class PosePredictor(DetectionPredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initializes PosePredictor, sets task to 'pose' and logs a warning for using 'mps' as device."""
+        """Initialize PosePredictor, set task to 'pose' and log a warning for using 'mps' as device."""
         super().__init__(cfg, overrides, _callbacks)
         self.args.task = "pose"
         if isinstance(self.args.device, str) and self.args.device.lower() == "mps":
@@ -28,19 +38,25 @@ class PosePredictor(DetectionPredictor):
 
     def construct_result(self, pred, img, orig_img, img_path):
         """
-        Constructs the result object from the prediction.
+        Construct the result object from the prediction, including keypoints.
+
+        This method extends the parent class implementation by extracting keypoint data from predictions
+        and adding them to the result object.
 
         Args:
-            pred (torch.Tensor): The predicted bounding boxes, scores, and keypoints.
-            img (torch.Tensor): The image after preprocessing.
-            orig_img (np.ndarray): The original image before preprocessing.
-            img_path (str): The path to the original image.
+            pred (torch.Tensor): The predicted bounding boxes, scores, and keypoints with shape (N, 6+K*D) where N is
+                the number of detections, K is the number of keypoints, and D is the keypoint dimension.
+            img (torch.Tensor): The processed input image tensor with shape (B, C, H, W).
+            orig_img (np.ndarray): The original unprocessed image as a numpy array.
+            img_path (str): The path to the original image file.
 
         Returns:
             (Results): The result object containing the original image, image path, class names, bounding boxes, and keypoints.
         """
         result = super().construct_result(pred, img, orig_img, img_path)
+        # Extract keypoints from prediction and reshape according to model's keypoint shape
         pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:]
+        # Scale keypoints coordinates to match the original image dimensions
         pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
         result.update(keypoints=pred_kpts)
         return result

ultralytics/models/yolo/pose/train.py

@@ -10,7 +10,23 @@ from ultralytics.utils.plotting import plot_images, plot_results
 
 class PoseTrainer(yolo.detect.DetectionTrainer):
     """
-    A class extending the DetectionTrainer class for training based on a pose model.
+    A class extending the DetectionTrainer class for training YOLO pose estimation models.
+
+    This trainer specializes in handling pose estimation tasks, managing model training, validation, and visualization
+    of pose keypoints alongside bounding boxes.
+
+    Attributes:
+        args (Dict): Configuration arguments for training.
+        model (PoseModel): The pose estimation model being trained.
+        data (Dict): Dataset configuration including keypoint shape information.
+        loss_names (Tuple[str]): Names of the loss components used in training.
+
+    Methods:
+        get_model: Retrieves a pose estimation model with specified configuration.
+        set_model_attributes: Sets keypoints shape attribute on the model.
+        get_validator: Creates a validator instance for model evaluation.
+        plot_training_samples: Visualizes training samples with keypoints.
+        plot_metrics: Generates and saves training/validation metric plots.
 
     Examples:
         >>> from ultralytics.models.yolo.pose import PoseTrainer

ultralytics/models/yolo/pose/val.py

@@ -16,6 +16,29 @@ class PoseValidator(DetectionValidator):
     """
     A class extending the DetectionValidator class for validation based on a pose model.
 
+    This validator is specifically designed for pose estimation tasks, handling keypoints and implementing
+    specialized metrics for pose evaluation.
+
+    Attributes:
+        sigma (np.ndarray): Sigma values for OKS calculation, either from OKS_SIGMA or ones divided by number of keypoints.
+        kpt_shape (List[int]): Shape of the keypoints, typically [17, 3] for COCO format.
+        args (Dict): Arguments for the validator including task set to "pose".
+        metrics (PoseMetrics): Metrics object for pose evaluation.
+
+    Methods:
+        preprocess: Preprocesses batch data for pose validation.
+        get_desc: Returns description of evaluation metrics.
+        init_metrics: Initializes pose metrics for the model.
+        _prepare_batch: Prepares a batch for processing.
+        _prepare_pred: Prepares and scales predictions for evaluation.
+        update_metrics: Updates metrics with new predictions.
+        _process_batch: Processes batch to compute IoU between detections and ground truth.
+        plot_val_samples: Plots validation samples with ground truth annotations.
+        plot_predictions: Plots model predictions.
+        save_one_txt: Saves detections to a text file.
+        pred_to_json: Converts predictions to COCO JSON format.
+        eval_json: Evaluates model using COCO JSON format.
+
     Examples:
         >>> from ultralytics.models.yolo.pose import PoseValidator
         >>> args = dict(model="yolo11n-pose.pt", data="coco8-pose.yaml")
@@ -24,7 +47,7 @@ class PoseValidator(DetectionValidator):
     """
 
     def __init__(self, dataloader=None, save_dir=None, pbar=None, args=None, _callbacks=None):
-        """Initialize a 'PoseValidator' object with custom parameters and assigned attributes."""
+        """Initialize a PoseValidator object with custom parameters and assigned attributes."""
         super().__init__(dataloader, save_dir, pbar, args, _callbacks)
         self.sigma = None
         self.kpt_shape = None
@@ -37,13 +60,13 @@ class PoseValidator(DetectionValidator):
             )
 
     def preprocess(self, batch):
-        """Preprocesses the batch by converting the 'keypoints' data into a float and moving it to the device."""
+        """Preprocess batch by converting keypoints data to float and moving it to the device."""
         batch = super().preprocess(batch)
         batch["keypoints"] = batch["keypoints"].to(self.device).float()
         return batch
 
     def get_desc(self):
-        """Returns description of evaluation metrics in string format."""
+        """Return description of evaluation metrics in string format."""
         return ("%22s" + "%11s" * 10) % (
             "Class",
             "Images",
@@ -59,7 +82,7 @@ class PoseValidator(DetectionValidator):
         )
 
     def init_metrics(self, model):
-        """Initiate pose estimation metrics for YOLO model."""
+        """Initialize pose estimation metrics for YOLO model."""
         super().init_metrics(model)
         self.kpt_shape = self.data["kpt_shape"]
         is_pose = self.kpt_shape == [17, 3]
@@ -68,7 +91,7 @@ class PoseValidator(DetectionValidator):
         self.stats = dict(tp_p=[], tp=[], conf=[], pred_cls=[], target_cls=[], target_img=[])
 
     def _prepare_batch(self, si, batch):
-        """Prepares a batch for processing by converting keypoints to float and moving to device."""
+        """Prepare a batch for processing by converting keypoints to float and scaling to original dimensions."""
         pbatch = super()._prepare_batch(si, batch)
         kpts = batch["keypoints"][batch["batch_idx"] == si]
         h, w = pbatch["imgsz"]
@@ -80,7 +103,7 @@ class PoseValidator(DetectionValidator):
         return pbatch
 
     def _prepare_pred(self, pred, pbatch):
-        """Prepares and scales keypoints in a batch for pose processing."""
+        """Prepare and scale keypoints in predictions for pose processing."""
         predn = super()._prepare_pred(pred, pbatch)
         nk = pbatch["kpts"].shape[1]
         pred_kpts = predn[:, 6:].view(len(predn), nk, -1)
@@ -88,7 +111,16 @@ class PoseValidator(DetectionValidator):
         return predn, pred_kpts
 
     def update_metrics(self, preds, batch):
-        """Metrics."""
+        """
+        Update metrics with new predictions and ground truth data.
+
+        This method processes each prediction, compares it with ground truth, and updates various statistics
+        for performance evaluation.
+
+        Args:
+            preds (List[torch.Tensor]): List of prediction tensors from the model.
+            batch (Dict): Batch data containing images and ground truth annotations.
+        """
         for si, pred in enumerate(preds):
             self.seen += 1
             npr = len(pred)
@@ -158,16 +190,9 @@ class PoseValidator(DetectionValidator):
             (torch.Tensor): A tensor with shape (N, 10) representing the correct prediction matrix for 10 IoU levels,
                 where N is the number of detections.
 
-        Examples:
-            >>> detections = torch.rand(100, 6)  # 100 predictions: (x1, y1, x2, y2, conf, class)
-            >>> gt_bboxes = torch.rand(50, 4)  # 50 ground truth boxes: (x1, y1, x2, y2)
-            >>> gt_cls = torch.randint(0, 2, (50,))  # 50 ground truth class indices
-            >>> pred_kpts = torch.rand(100, 51)  # 100 predicted keypoints
-            >>> gt_kpts = torch.rand(50, 51)  # 50 ground truth keypoints
-            >>> correct_preds = _process_batch(detections, gt_bboxes, gt_cls, pred_kpts, gt_kpts)
-
-        Note:
-            `0.53` scale factor used in area computation is referenced from https://github.com/jin-s13/xtcocoapi/blob/master/xtcocotools/cocoeval.py#L384.
+        Notes:
+            `0.53` scale factor used in area computation is referenced from
+            https://github.com/jin-s13/xtcocoapi/blob/master/xtcocotools/cocoeval.py#L384.
         """
         if pred_kpts is not None and gt_kpts is not None:
             # `0.53` is from https://github.com/jin-s13/xtcocoapi/blob/master/xtcocotools/cocoeval.py#L384
@@ -179,7 +204,7 @@ class PoseValidator(DetectionValidator):
         return self.match_predictions(detections[:, 5], gt_cls, iou)
 
     def plot_val_samples(self, batch, ni):
-        """Plots and saves validation set samples with predicted bounding boxes and keypoints."""
+        """Plot and save validation set samples with ground truth bounding boxes and keypoints."""
         plot_images(
             batch["img"],
             batch["batch_idx"],
@@ -193,7 +218,7 @@ class PoseValidator(DetectionValidator):
         )
 
     def plot_predictions(self, batch, preds, ni):
-        """Plots predictions for YOLO model."""
+        """Plot and save model predictions with bounding boxes and keypoints."""
         pred_kpts = torch.cat([p[:, 6:].view(-1, *self.kpt_shape) for p in preds], 0)
         plot_images(
             batch["img"],
@@ -218,7 +243,7 @@ class PoseValidator(DetectionValidator):
         ).save_txt(file, save_conf=save_conf)
 
     def pred_to_json(self, predn, filename):
-        """Converts YOLO predictions to COCO JSON format."""
+        """Convert YOLO predictions to COCO JSON format."""
         stem = Path(filename).stem
         image_id = int(stem) if stem.isnumeric() else stem
         box = ops.xyxy2xywh(predn[:, :4])  # xywh
@@ -235,7 +260,7 @@ class PoseValidator(DetectionValidator):
             )
 
     def eval_json(self, stats):
-        """Evaluates object detection model using COCO JSON format."""
+        """Evaluate object detection model using COCO JSON format."""
         if self.args.save_json and self.is_coco and len(self.jdict):
             anno_json = self.data["path"] / "annotations/person_keypoints_val2017.json"  # annotations
             pred_json = self.save_dir / "predictions.json"  # predictions

ultralytics/models/yolo/segment/predict.py

@@ -9,6 +9,19 @@ class SegmentationPredictor(DetectionPredictor):
     """
     A class extending the DetectionPredictor class for prediction based on a segmentation model.
 
+    This class specializes in processing segmentation model outputs, handling both bounding boxes and masks in the
+    prediction results.
+
+    Attributes:
+        args (Dict): Configuration arguments for the predictor.
+        model (torch.nn.Module): The loaded YOLO segmentation model.
+        batch (List): Current batch of images being processed.
+
+    Methods:
+        postprocess: Applies non-max suppression and processes detections.
+        construct_results: Constructs a list of result objects from predictions.
+        construct_result: Constructs a single result object from a prediction.
+
     Examples:
         >>> from ultralytics.utils import ASSETS
         >>> from ultralytics.models.yolo.segment import SegmentationPredictor
@@ -18,19 +31,19 @@ class SegmentationPredictor(DetectionPredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """Initializes the SegmentationPredictor with the provided configuration, overrides, and callbacks."""
+        """Initialize the SegmentationPredictor with configuration, overrides, and callbacks."""
         super().__init__(cfg, overrides, _callbacks)
         self.args.task = "segment"
 
     def postprocess(self, preds, img, orig_imgs):
-        """Applies non-max suppression and processes detections for each image in an input batch."""
-        # tuple if PyTorch model or array if exported
+        """Apply non-max suppression and process detections for each image in the input batch."""
+        # Extract protos - tuple if PyTorch model or array if exported
         protos = preds[1][-1] if isinstance(preds[1], tuple) else preds[1]
         return super().postprocess(preds[0], img, orig_imgs, protos=protos)
 
     def construct_results(self, preds, img, orig_imgs, protos):
         """
-        Constructs a list of result objects from the predictions.
+        Construct a list of result objects from the predictions.
 
         Args:
             preds (List[torch.Tensor]): List of predicted bounding boxes, scores, and masks.
@@ -39,7 +52,8 @@ class SegmentationPredictor(DetectionPredictor):
             protos (List[torch.Tensor]): List of prototype masks.
 
         Returns:
-            (list): List of result objects containing the original images, image paths, class names, bounding boxes, and masks.
+            (List[Results]): List of result objects containing the original images, image paths, class names,
+                bounding boxes, and masks.
         """
         return [
             self.construct_result(pred, img, orig_img, img_path, proto)
@@ -48,7 +62,7 @@ class SegmentationPredictor(DetectionPredictor):
 
     def construct_result(self, pred, img, orig_img, img_path, proto):
         """
-        Constructs the result object from the prediction.
+        Construct a single result object from the prediction.
 
         Args:
             pred (np.ndarray): The predicted bounding boxes, scores, and masks.
@@ -58,7 +72,7 @@ class SegmentationPredictor(DetectionPredictor):
             proto (torch.Tensor): The prototype masks.
 
         Returns:
-            (Results): The result object containing the original image, image path, class names, bounding boxes, and masks.
+            (Results): Result object containing the original image, image path, class names, bounding boxes, and masks.
         """
         if not len(pred):  # save empty boxes
             masks = None
@@ -69,6 +83,6 @@ class SegmentationPredictor(DetectionPredictor):
             masks = ops.process_mask(proto, pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True)  # HWC
             pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
         if masks is not None:
-            keep = masks.sum((-2, -1)) > 0  # only keep preds with masks
+            keep = masks.sum((-2, -1)) > 0  # only keep predictions with masks
             pred, masks = pred[keep], masks[keep]
         return Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks)

ultralytics/models/yolo/segment/train.py

@@ -12,6 +12,12 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer):
     """
     A class extending the DetectionTrainer class for training based on a segmentation model.
 
+    This trainer specializes in handling segmentation tasks, extending the detection trainer with segmentation-specific
+    functionality including model initialization, validation, and visualization.
+
+    Attributes:
+        loss_names (Tuple[str]): Names of the loss components used during training.
+
     Examples:
         >>> from ultralytics.models.yolo.segment import SegmentationTrainer
         >>> args = dict(model="yolo11n-seg.pt", data="coco8-seg.yaml", epochs=3)