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

ultralytics/models/yolo/pose/val.py

@@ -16,18 +16,38 @@ class PoseValidator(DetectionValidator):
     """
     A class extending the DetectionValidator class for validation based on a pose model.
 
-    Example:
-        ```python
-        from ultralytics.models.yolo.pose import PoseValidator
-
-        args = dict(model="yolo11n-pose.pt", data="coco8-pose.yaml")
-        validator = PoseValidator(args=args)
-        validator()
-        ```
+    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")
+        >>> validator = PoseValidator(args=args)
+        >>> validator()
     """
 
     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
@@ -40,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",
@@ -62,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]
@@ -71,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"]
@@ -83,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)
@@ -91,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)
@@ -161,18 +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.
 
-        Example:
-            ```python
-            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
@@ -184,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"],
@@ -198,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"],
@@ -223,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
@@ -240,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,31 +9,41 @@ class SegmentationPredictor(DetectionPredictor):
     """
     A class extending the DetectionPredictor class for prediction based on a segmentation model.
 
-    Example:
-        ```python
-        from ultralytics.utils import ASSETS
-        from ultralytics.models.yolo.segment import SegmentationPredictor
+    This class specializes in processing segmentation model outputs, handling both bounding boxes and masks in the
+    prediction results.
 
-        args = dict(model="yolo11n-seg.pt", source=ASSETS)
-        predictor = SegmentationPredictor(overrides=args)
-        predictor.predict_cli()
-        ```
+    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
+        >>> args = dict(model="yolo11n-seg.pt", source=ASSETS)
+        >>> predictor = SegmentationPredictor(overrides=args)
+        >>> predictor.predict_cli()
     """
 
     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.
@@ -42,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)
@@ -51,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.
@@ -61,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
@@ -72,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,14 +12,17 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer):
     """
     A class extending the DetectionTrainer class for training based on a segmentation model.
 
-    Example:
-        ```python
-        from ultralytics.models.yolo.segment import SegmentationTrainer
-
-        args = dict(model="yolo11n-seg.pt", data="coco8-seg.yaml", epochs=3)
-        trainer = SegmentationTrainer(overrides=args)
-        trainer.train()
-        ```
+    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)
+        >>> trainer = SegmentationTrainer(overrides=args)
+        >>> trainer.train()
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):

ultralytics/models/yolo/segment/val.py

@@ -18,18 +18,34 @@ class SegmentationValidator(DetectionValidator):
     """
     A class extending the DetectionValidator class for validation based on a segmentation model.
 
-    Example:
-        ```python
-        from ultralytics.models.yolo.segment import SegmentationValidator
-
-        args = dict(model="yolo11n-seg.pt", data="coco8-seg.yaml")
-        validator = SegmentationValidator(args=args)
-        validator()
-        ```
+    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")
+        >>> validator = SegmentationValidator(args=args)
+        >>> validator()
     """
 
     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
@@ -37,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:
@@ -69,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)
@@ -157,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
 
@@ -171,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.
 
@@ -184,13 +240,11 @@ class SegmentationValidator(DetectionValidator):
             - If `masks` is True, the function computes IoU between predicted and ground truth masks.
             - If `overlap` is True and `masks` is True, overlapping masks are taken into account when computing IoU.
 
-        Example:
-            ```python
-            detections = torch.tensor([[25, 30, 200, 300, 0.8, 1], [50, 60, 180, 290, 0.75, 0]])
-            gt_bboxes = torch.tensor([[24, 29, 199, 299], [55, 65, 185, 295]])
-            gt_cls = torch.tensor([1, 0])
-            correct_preds = validator._process_batch(detections, gt_bboxes, gt_cls)
-            ```
+        Examples:
+            >>> detections = torch.tensor([[25, 30, 200, 300, 0.8, 1], [50, 60, 180, 290, 0.75, 0]])
+            >>> gt_bboxes = torch.tensor([[24, 29, 199, 299], [55, 65, 185, 295]])
+            >>> gt_cls = torch.tensor([1, 0])
+            >>> correct_preds = validator._process_batch(detections, gt_bboxes, gt_cls)
         """
         if masks:
             if overlap:
@@ -208,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"],
@@ -222,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
@@ -235,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(
@@ -248,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,18 +25,32 @@ class WorldTrainer(yolo.detect.DetectionTrainer):
     """
     A class to fine-tune a world model on a close-set dataset.
 
-    Example:
-        ```python
-        from ultralytics.models.yolo.world import WorldModel
-
-        args = dict(model="yolov8s-world.pt", data="coco8.yaml", epochs=3)
-        trainer = WorldTrainer(overrides=args)
-        trainer.train()
-        ```
+    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)
+        >>> trainer = WorldTrainer(overrides=args)
+        >>> trainer.train()
     """
 
     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)
@@ -50,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(
@@ -67,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(
@@ -80,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