dgenerate-ultralytics-headless 8.3.143__py3-none-any.whl → 8.3.144__py3-none-any.whl

ultralytics/models/yolo/model.py

@@ -1,6 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
 
 from ultralytics.data.build import load_inference_source
 from ultralytics.engine.model import Model
@@ -19,9 +20,34 @@ from ultralytics.utils import ROOT, YAML
 
 
 class YOLO(Model):
-    """YOLO (You Only Look Once) object detection model."""
+    """
+    YOLO (You Only Look Once) object detection model.
 
-    def __init__(self, model="yolo11n.pt", task=None, verbose=False):
+    This class provides a unified interface for YOLO models, automatically switching to specialized model types
+    (YOLOWorld or YOLOE) based on the model filename. It supports various computer vision tasks including object
+    detection, segmentation, classification, pose estimation, and oriented bounding box detection.
+
+    Attributes:
+        model: The loaded YOLO model instance.
+        task: The task type (detect, segment, classify, pose, obb).
+        overrides: Configuration overrides for the model.
+
+    Methods:
+        __init__: Initialize a YOLO model with automatic type detection.
+        task_map: Map tasks to their corresponding model, trainer, validator, and predictor classes.
+
+    Examples:
+        Load a pretrained YOLOv11n detection model
+        >>> model = YOLO("yolo11n.pt")
+
+        Load a pretrained YOLO11n segmentation model
+        >>> model = YOLO("yolo11n-seg.pt")
+
+        Initialize from a YAML configuration
+        >>> model = YOLO("yolo11n.yaml")
+    """
+
+    def __init__(self, model: Union[str, Path] = "yolo11n.pt", task: Optional[str] = None, verbose: bool = False):
         """
         Initialize a YOLO model.
 
@@ -30,7 +56,7 @@ class YOLO(Model):
 
         Args:
             model (str | Path): Model name or path to model file, i.e. 'yolo11n.pt', 'yolo11n.yaml'.
-            task (str | None): YOLO task specification, i.e. 'detect', 'segment', 'classify', 'pose', 'obb'.
+            task (str, optional): YOLO task specification, i.e. 'detect', 'segment', 'classify', 'pose', 'obb'.
                 Defaults to auto-detection based on model.
             verbose (bool): Display model info on load.
 
@@ -59,7 +85,7 @@ class YOLO(Model):
                 self.__dict__ = new_instance.__dict__
 
     @property
-    def task_map(self):
+    def task_map(self) -> Dict[str, Dict[str, Any]]:
         """Map head to model, trainer, validator, and predictor classes."""
         return {
             "classify": {
@@ -96,9 +122,32 @@ class YOLO(Model):
 
 
 class YOLOWorld(Model):
-    """YOLO-World object detection model."""
+    """
+    YOLO-World object detection model.
 
-    def __init__(self, model="yolov8s-world.pt", verbose=False) -> None:
+    YOLO-World is an open-vocabulary object detection model that can detect objects based on text descriptions
+    without requiring training on specific classes. It extends the YOLO architecture to support real-time
+    open-vocabulary detection.
+
+    Attributes:
+        model: The loaded YOLO-World model instance.
+        task: Always set to 'detect' for object detection.
+        overrides: Configuration overrides for the model.
+
+    Methods:
+        __init__: Initialize YOLOv8-World model with a pre-trained model file.
+        task_map: Map tasks to their corresponding model, trainer, validator, and predictor classes.
+        set_classes: Set the model's class names for detection.
+
+    Examples:
+        Load a YOLOv8-World model
+        >>> model = YOLOWorld("yolov8s-world.pt")
+
+        Set custom classes for detection
+        >>> model.set_classes(["person", "car", "bicycle"])
+    """
+
+    def __init__(self, model: Union[str, Path] = "yolov8s-world.pt", verbose: bool = False) -> None:
         """
         Initialize YOLOv8-World model with a pre-trained model file.
 
@@ -116,7 +165,7 @@ class YOLOWorld(Model):
             self.model.names = YAML.load(ROOT / "cfg/datasets/coco8.yaml").get("names")
 
     @property
-    def task_map(self):
+    def task_map(self) -> Dict[str, Dict[str, Any]]:
         """Map head to model, validator, and predictor classes."""
         return {
             "detect": {
@@ -127,12 +176,12 @@ class YOLOWorld(Model):
             }
         }
 
-    def set_classes(self, classes):
+    def set_classes(self, classes: List[str]) -> None:
         """
         Set the model's class names for detection.
 
         Args:
-            classes (list[str]): A list of categories i.e. ["person"].
+            classes (List[str]): A list of categories i.e. ["person"].
         """
         self.model.set_classes(classes)
         # Remove background if it's given
@@ -147,9 +196,43 @@ class YOLOWorld(Model):
 
 
 class YOLOE(Model):
-    """YOLOE object detection and segmentation model."""
-
-    def __init__(self, model="yoloe-11s-seg.pt", task=None, verbose=False) -> None:
+    """
+    YOLOE object detection and segmentation model.
+
+    YOLOE is an enhanced YOLO model that supports both object detection and instance segmentation tasks with
+    improved performance and additional features like visual and text positional embeddings.
+
+    Attributes:
+        model: The loaded YOLOE model instance.
+        task: The task type (detect or segment).
+        overrides: Configuration overrides for the model.
+
+    Methods:
+        __init__: Initialize YOLOE model with a pre-trained model file.
+        task_map: Map tasks to their corresponding model, trainer, validator, and predictor classes.
+        get_text_pe: Get text positional embeddings for the given texts.
+        get_visual_pe: Get visual positional embeddings for the given image and visual features.
+        set_vocab: Set vocabulary and class names for the YOLOE model.
+        get_vocab: Get vocabulary for the given class names.
+        set_classes: Set the model's class names and embeddings for detection.
+        val: Validate the model using text or visual prompts.
+        predict: Run prediction on images, videos, directories, streams, etc.
+
+    Examples:
+        Load a YOLOE detection model
+        >>> model = YOLOE("yoloe-11s-seg.pt")
+
+        Set vocabulary and class names
+        >>> model.set_vocab(["person", "car", "dog"], ["person", "car", "dog"])
+
+        Predict with visual prompts
+        >>> prompts = {"bboxes": [[10, 20, 100, 200]], "cls": ["person"]}
+        >>> results = model.predict("image.jpg", visual_prompts=prompts)
+    """
+
+    def __init__(
+        self, model: Union[str, Path] = "yoloe-11s-seg.pt", task: Optional[str] = None, verbose: bool = False
+    ) -> None:
         """
         Initialize YOLOE model with a pre-trained model file.
 
@@ -165,7 +248,7 @@ class YOLOE(Model):
             self.model.names = YAML.load(ROOT / "cfg/datasets/coco8.yaml").get("names")
 
     @property
-    def task_map(self):
+    def task_map(self) -> Dict[str, Dict[str, Any]]:
         """Map head to model, validator, and predictor classes."""
         return {
             "detect": {
@@ -210,7 +293,7 @@ class YOLOE(Model):
         assert isinstance(self.model, YOLOEModel)
         return self.model.get_visual_pe(img, visual)
 
-    def set_vocab(self, vocab, names):
+    def set_vocab(self, vocab: List[str], names: List[str]) -> None:
         """
         Set vocabulary and class names for the YOLOE model.
 
@@ -218,8 +301,8 @@ class YOLOE(Model):
         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.
+            vocab (List[str]): Vocabulary list containing tokens or words used by the model for text processing.
+            names (List[str]): List of class names that the model can detect or classify.
 
         Raises:
             AssertionError: If the model is not an instance of YOLOEModel.
@@ -236,12 +319,12 @@ class YOLOE(Model):
         assert isinstance(self.model, YOLOEModel)
         return self.model.get_vocab(names)
 
-    def set_classes(self, classes, embeddings):
+    def set_classes(self, classes: List[str], embeddings) -> None:
         """
         Set the model's class names and embeddings for detection.
 
         Args:
-            classes (list[str]): A list of categories i.e. ["person"].
+            classes (List[str]): A list of categories i.e. ["person"].
             embeddings (torch.Tensor): Embeddings corresponding to the classes.
         """
         assert isinstance(self.model, YOLOEModel)
@@ -257,8 +340,8 @@ class YOLOE(Model):
     def val(
         self,
         validator=None,
-        load_vp=False,
-        refer_data=None,
+        load_vp: bool = False,
+        refer_data: Optional[str] = None,
         **kwargs,
     ):
         """
@@ -285,7 +368,7 @@ class YOLOE(Model):
         self,
         source=None,
         stream: bool = False,
-        visual_prompts: dict = {},
+        visual_prompts: Dict[str, List] = {},
         refer_image=None,
         predictor=None,
         **kwargs,
@@ -298,8 +381,8 @@ class YOLOE(Model):
                 directory paths, URL/YouTube streams, PIL images, numpy arrays, or webcam indices.
             stream (bool): Whether to stream the prediction results. If True, results are yielded as a
                 generator as they are computed.
-            visual_prompts (dict): Dictionary containing visual prompts for the model. Must include 'bboxes' and
-                'cls' keys when non-empty.
+            visual_prompts (Dict[str, List]): Dictionary containing visual prompts for the model. Must include
+                'bboxes' and 'cls' keys when non-empty.
             refer_image (str | PIL.Image | np.ndarray, optional): Reference image for visual prompts.
             predictor (callable, optional): Custom predictor function. If None, a predictor is automatically
                 loaded based on the task.

ultralytics/models/yolo/obb/predict.py

@@ -30,8 +30,6 @@ class OBBPredictor(DetectionPredictor):
         """
         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.
@@ -51,14 +49,15 @@ class OBBPredictor(DetectionPredictor):
         Construct the result object from the prediction.
 
         Args:
-            pred (torch.Tensor): The predicted bounding boxes, scores, and rotation angles with shape (N, 6) where
+            pred (torch.Tensor): The predicted bounding boxes, scores, and rotation angles with shape (N, 7) 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.
 
         Returns:
-            (Results): The result object containing the original image, image path, class names, and oriented bounding boxes.
+            (Results): The result object containing the original image, image path, class names, and oriented bounding
+                boxes.
         """
         rboxes = ops.regularize_rboxes(torch.cat([pred[:, :4], pred[:, -1:]], dim=-1))
         rboxes[:, :4] = ops.scale_boxes(img.shape[2:], rboxes[:, :4], orig_img.shape, xywh=True)

ultralytics/models/yolo/obb/train.py

@@ -1,6 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from copy import copy
+from pathlib import Path
+from typing import Any, List, Optional, Union
 
 from ultralytics.models import yolo
 from ultralytics.nn.tasks import OBBModel
@@ -11,8 +13,12 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
     """
     A class extending the DetectionTrainer class for training based on an Oriented Bounding Box (OBB) model.
 
+    This trainer specializes in training YOLO models that detect oriented bounding boxes, which are useful for
+    detecting objects at arbitrary angles rather than just axis-aligned rectangles.
+
     Attributes:
-        loss_names (Tuple[str]): Names of the loss components used during training.
+        loss_names (tuple): Names of the loss components used during training including box_loss, cls_loss,
+            and dfl_loss.
 
     Methods:
         get_model: Return OBBModel initialized with specified config and weights.
@@ -25,7 +31,7 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
         >>> trainer.train()
     """
 
-    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
+    def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[dict] = None, _callbacks: Optional[List[Any]] = None):
         """
         Initialize an OBBTrainer object for training Oriented Bounding Box (OBB) models.
 
@@ -37,7 +43,7 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
                 model configuration.
             overrides (dict, optional): Dictionary of parameter overrides for the configuration. Any values here
                 will take precedence over those in cfg.
-            _callbacks (list, optional): List of callback functions to be invoked during training.
+            _callbacks (List[Any], optional): List of callback functions to be invoked during training.
 
         Examples:
             >>> from ultralytics.models.yolo.obb import OBBTrainer
@@ -50,14 +56,16 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
         overrides["task"] = "obb"
         super().__init__(cfg, overrides, _callbacks)
 
-    def get_model(self, cfg=None, weights=None, verbose=True):
+    def get_model(
+        self, cfg: Optional[Union[str, dict]] = None, weights: Optional[Union[str, Path]] = None, verbose: bool = True
+    ) -> OBBModel:
         """
         Return OBBModel initialized with specified config and weights.
 
         Args:
-            cfg (str | dict | None): Model configuration. Can be a path to a YAML config file, a dictionary
+            cfg (str | dict, optional): Model configuration. Can be a path to a YAML config file, a dictionary
                 containing configuration parameters, or None to use default configuration.
-            weights (str | Path | None): Path to pretrained weights file. If None, random initialization is used.
+            weights (str | Path, optional): Path to pretrained weights file. If None, random initialization is used.
             verbose (bool): Whether to display model information during initialization.
 
         Returns:

ultralytics/models/yolo/obb/val.py

@@ -1,6 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from pathlib import Path
+from typing import Dict, List, Tuple, Union
 
 import torch
 
@@ -63,34 +64,31 @@ class OBBValidator(DetectionValidator):
         val = self.data.get(self.args.split, "")  # validation path
         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):
+    def _process_batch(self, detections: torch.Tensor, gt_bboxes: torch.Tensor, gt_cls: torch.Tensor) -> torch.Tensor:
         """
-        Perform computation of the correct prediction matrix for a batch of detections and ground truth bounding boxes.
+        Compute the correct prediction matrix for a batch of detections and ground truth bounding boxes.
 
         Args:
-            detections (torch.Tensor): A tensor of shape (N, 7) representing the detected bounding boxes and associated
-                data. Each detection is represented as (x1, y1, x2, y2, conf, class, angle).
-            gt_bboxes (torch.Tensor): A tensor of shape (M, 5) representing the ground truth bounding boxes. Each box is
-                represented as (x1, y1, x2, y2, angle).
-            gt_cls (torch.Tensor): A tensor of shape (M,) representing class labels for the ground truth bounding boxes.
+            detections (torch.Tensor): Detected bounding boxes and associated data with shape (N, 7) where each
+                detection is represented as (x1, y1, x2, y2, conf, class, angle).
+            gt_bboxes (torch.Tensor): Ground truth bounding boxes with shape (M, 5) where each box is represented
+                as (x1, y1, x2, y2, angle).
+            gt_cls (torch.Tensor): Class labels for the ground truth bounding boxes with shape (M,).
 
         Returns:
-            (torch.Tensor): The correct prediction matrix with shape (N, 10), which includes 10 IoU (Intersection over
-                Union) levels for each detection, indicating the accuracy of predictions compared to the ground truth.
+            (torch.Tensor): The correct prediction matrix with shape (N, 10), which includes 10 IoU levels for each
+                detection, indicating the accuracy of predictions compared to the ground truth.
 
         Examples:
             >>> detections = torch.rand(100, 7)  # 100 sample detections
             >>> gt_bboxes = torch.rand(50, 5)  # 50 sample ground truth boxes
             >>> gt_cls = torch.randint(0, 5, (50,))  # 50 ground truth class labels
-            >>> correct_matrix = OBBValidator._process_batch(detections, gt_bboxes, gt_cls)
-
-        Note:
-            This method relies on `batch_probiou` to calculate IoU between detections and ground truth bounding boxes.
+            >>> correct_matrix = validator._process_batch(detections, gt_bboxes, gt_cls)
         """
         iou = batch_probiou(gt_bboxes, torch.cat([detections[:, :4], detections[:, -1:]], dim=-1))
         return self.match_predictions(detections[:, 5], gt_cls, iou)
 
-    def _prepare_batch(self, si, batch):
+    def _prepare_batch(self, si: int, batch: Dict) -> Dict:
         """
         Prepare batch data for OBB validation with proper scaling and formatting.
 
@@ -104,8 +102,8 @@ class OBBValidator(DetectionValidator):
                 - img: Batch of images
                 - ratio_pad: Ratio and padding information
 
-        This method filters the batch data for a specific batch index, extracts class labels and bounding boxes,
-        and scales the bounding boxes to the original image dimensions.
+        Returns:
+            (dict): Prepared batch data with scaled bounding boxes and metadata.
         """
         idx = batch["batch_idx"] == si
         cls = batch["cls"][idx].squeeze(-1)
@@ -118,7 +116,7 @@ class OBBValidator(DetectionValidator):
             ops.scale_boxes(imgsz, bbox, ori_shape, ratio_pad=ratio_pad, xywh=True)  # native-space labels
         return {"cls": cls, "bbox": bbox, "ori_shape": ori_shape, "imgsz": imgsz, "ratio_pad": ratio_pad}
 
-    def _prepare_pred(self, pred, pbatch):
+    def _prepare_pred(self, pred: torch.Tensor, pbatch: Dict) -> torch.Tensor:
         """
         Prepare predictions by scaling bounding boxes to original image dimensions.
 
@@ -141,7 +139,7 @@ class OBBValidator(DetectionValidator):
         )  # native-space pred
         return predn
 
-    def plot_predictions(self, batch, preds, ni):
+    def plot_predictions(self, batch: Dict, preds: List[torch.Tensor], ni: int):
         """
         Plot predicted bounding boxes on input images and save the result.
 
@@ -165,7 +163,7 @@ class OBBValidator(DetectionValidator):
             on_plot=self.on_plot,
         )  # pred
 
-    def pred_to_json(self, predn, filename):
+    def pred_to_json(self, predn: torch.Tensor, filename: Union[str, Path]):
         """
         Convert YOLO predictions to COCO JSON format with rotated bounding box information.
 
@@ -194,9 +192,9 @@ class OBBValidator(DetectionValidator):
                 }
             )
 
-    def save_one_txt(self, predn, save_conf, shape, file):
+    def save_one_txt(self, predn: torch.Tensor, save_conf: bool, shape: Tuple[int, int], file: Union[Path, str]):
         """
-        Save YOLO OBB (Oriented Bounding Box) detections to a text file in normalized coordinates.
+        Save YOLO OBB detections to a text file in normalized coordinates.
 
         Args:
             predn (torch.Tensor): Predicted detections with shape (N, 7) containing bounding boxes, confidence scores,
@@ -224,8 +222,16 @@ class OBBValidator(DetectionValidator):
             obb=obb,
         ).save_txt(file, save_conf=save_conf)
 
-    def eval_json(self, stats):
-        """Evaluate YOLO output in JSON format and save predictions in DOTA format."""
+    def eval_json(self, stats: Dict) -> Dict:
+        """
+        Evaluate YOLO output in JSON format and save predictions in DOTA format.
+
+        Args:
+            stats (dict): Performance statistics dictionary.
+
+        Returns:
+            (dict): Updated performance statistics.
+        """
         if self.args.save_json and self.is_dota and len(self.jdict):
             import json
             import re

ultralytics/models/yolo/pose/predict.py

@@ -16,7 +16,7 @@ class PosePredictor(DetectionPredictor):
         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.
+        construct_result: Construct the result object from the prediction, including keypoints.
 
     Examples:
         >>> from ultralytics.utils import ASSETS
@@ -28,13 +28,13 @@ class PosePredictor(DetectionPredictor):
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
         """
-        Initialize PosePredictor, a specialized predictor for pose estimation tasks.
+        Initialize PosePredictor for pose estimation tasks.
 
-        This initializer sets up a PosePredictor instance, configuring it for pose detection tasks and handling
-        device-specific warnings for Apple MPS.
+        Sets up a PosePredictor instance, configuring it for pose detection tasks and handling device-specific
+        warnings for Apple MPS.
 
         Args:
-            cfg (Any): Configuration for the predictor. Default is DEFAULT_CFG.
+            cfg (Any): Configuration for the predictor.
             overrides (dict, optional): Configuration overrides that take precedence over cfg.
             _callbacks (list, optional): List of callback functions to be invoked during prediction.
 
@@ -57,8 +57,8 @@ class PosePredictor(DetectionPredictor):
         """
         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.
+        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 with shape (N, 6+K*D) where N is
@@ -68,7 +68,8 @@ class PosePredictor(DetectionPredictor):
             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.
+            (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

ultralytics/models/yolo/pose/train.py

@@ -1,6 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from copy import copy
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
 
 from ultralytics.models import yolo
 from ultralytics.nn.tasks import PoseModel
@@ -19,14 +21,15 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
         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.
+        loss_names (tuple): 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.
+        get_model: Retrieve a pose estimation model with specified configuration.
+        set_model_attributes: Set keypoints shape attribute on the model.
+        get_validator: Create a validator instance for model evaluation.
+        plot_training_samples: Visualize training samples with keypoints.
+        plot_metrics: Generate and save training/validation metric plots.
+        get_dataset: Retrieve the dataset and ensure it contains required kpt_shape key.
 
     Examples:
         >>> from ultralytics.models.yolo.pose import PoseTrainer
@@ -35,7 +38,7 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
         >>> trainer.train()
     """
 
-    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
+    def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[Dict[str, Any]] = None, _callbacks=None):
         """
         Initialize a PoseTrainer object for training YOLO pose estimation models.
 
@@ -68,13 +71,18 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
                 "See https://github.com/ultralytics/ultralytics/issues/4031."
             )
 
-    def get_model(self, cfg=None, weights=None, verbose=True):
+    def get_model(
+        self,
+        cfg: Optional[Union[str, Path, Dict[str, Any]]] = None,
+        weights: Optional[Union[str, Path]] = None,
+        verbose: bool = True,
+    ) -> PoseModel:
         """
         Get pose estimation model with specified configuration and weights.
 
         Args:
-            cfg (str | Path | dict | None): Model configuration file path or dictionary.
-            weights (str | Path | None): Path to the model weights file.
+            cfg (str | Path | dict, optional): Model configuration file path or dictionary.
+            weights (str | Path, optional): Path to the model weights file.
             verbose (bool): Whether to display model information.
 
         Returns:
@@ -89,18 +97,18 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
         return model
 
     def set_model_attributes(self):
-        """Sets keypoints shape attribute of PoseModel."""
+        """Set keypoints shape attribute of PoseModel."""
         super().set_model_attributes()
         self.model.kpt_shape = self.data["kpt_shape"]
 
     def get_validator(self):
-        """Returns an instance of the PoseValidator class for validation."""
+        """Return an instance of the PoseValidator class for validation."""
         self.loss_names = "box_loss", "pose_loss", "kobj_loss", "cls_loss", "dfl_loss"
         return yolo.pose.PoseValidator(
             self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
 
-    def plot_training_samples(self, batch, ni):
+    def plot_training_samples(self, batch: Dict[str, Any], ni: int):
         """
         Plot a batch of training samples with annotated class labels, bounding boxes, and keypoints.
 
@@ -135,12 +143,12 @@ class PoseTrainer(yolo.detect.DetectionTrainer):
         )
 
     def plot_metrics(self):
-        """Plots training/val metrics."""
+        """Plot training/validation metrics."""
         plot_results(file=self.csv, pose=True, on_plot=self.on_plot)  # save results.png
 
-    def get_dataset(self):
+    def get_dataset(self) -> Dict[str, Any]:
         """
-        Retrieves the dataset and ensures it contains the required `kpt_shape` key.
+        Retrieve the dataset and ensure it contains the required `kpt_shape` key.
 
         Returns:
             (dict): A dictionary containing the training/validation/test dataset and category names.