dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.3.248__py3-none-any.whl

ultralytics/models/yolo/segment/predict.py

@@ -6,8 +6,7 @@ from ultralytics.utils import DEFAULT_CFG, ops
 
 
 class SegmentationPredictor(DetectionPredictor):
-    """
-    A class extending the DetectionPredictor class for prediction based on a segmentation model.
+    """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.
@@ -31,8 +30,7 @@ class SegmentationPredictor(DetectionPredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """
-        Initialize the SegmentationPredictor with configuration, overrides, and callbacks.
+        """Initialize the SegmentationPredictor with configuration, overrides, and callbacks.
 
         This class specializes in processing segmentation model outputs, handling both bounding boxes and masks in the
         prediction results.
@@ -46,8 +44,7 @@ class SegmentationPredictor(DetectionPredictor):
         self.args.task = "segment"
 
     def postprocess(self, preds, img, orig_imgs):
-        """
-        Apply non-max suppression and process segmentation detections for each image in the input batch.
+        """Apply non-max suppression and process segmentation detections for each image in the input batch.
 
         Args:
             preds (tuple): Model predictions, containing bounding boxes, scores, classes, and mask coefficients.
@@ -55,8 +52,8 @@ class SegmentationPredictor(DetectionPredictor):
             orig_imgs (list | torch.Tensor | np.ndarray): Original image or batch of images.
 
         Returns:
-            (list): List of Results objects containing the segmentation predictions for each image in the batch.
-                Each Results object includes both bounding boxes and segmentation masks.
+            (list): List of Results objects containing the segmentation predictions for each image in the batch. Each
+                Results object includes both bounding boxes and segmentation masks.
 
         Examples:
             >>> predictor = SegmentationPredictor(overrides=dict(model="yolo11n-seg.pt"))
@@ -67,8 +64,7 @@ class SegmentationPredictor(DetectionPredictor):
         return super().postprocess(preds[0], img, orig_imgs, protos=protos)
 
     def construct_results(self, preds, img, orig_imgs, protos):
-        """
-        Construct 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.
@@ -77,8 +73,8 @@ class SegmentationPredictor(DetectionPredictor):
             protos (list[torch.Tensor]): List of prototype masks.
 
         Returns:
-            (list[Results]): 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)
@@ -86,8 +82,7 @@ class SegmentationPredictor(DetectionPredictor):
         ]
 
     def construct_result(self, pred, img, orig_img, img_path, proto):
-        """
-        Construct a single result object from the prediction.
+        """Construct a single result object from the prediction.
 
         Args:
             pred (torch.Tensor): The predicted bounding boxes, scores, and masks.
@@ -103,11 +98,12 @@ class SegmentationPredictor(DetectionPredictor):
             masks = None
         elif self.args.retina_masks:
             pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
-            masks = ops.process_mask_native(proto, pred[:, 6:], pred[:, :4], orig_img.shape[:2])  # HWC
+            masks = ops.process_mask_native(proto, pred[:, 6:], pred[:, :4], orig_img.shape[:2])  # NHW
         else:
-            masks = ops.process_mask(proto, pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True)  # HWC
+            masks = ops.process_mask(proto, pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True)  # NHW
             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 predictions with masks
-            pred, masks = pred[keep], masks[keep]
+            keep = masks.amax((-2, -1)) > 0  # only keep predictions with masks
+            if not all(keep):  # most predictions have masks
+                pred, masks = pred[keep], masks[keep]  # indexing is slow
         return Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks)

ultralytics/models/yolo/segment/train.py

@@ -11,8 +11,7 @@ from ultralytics.utils import DEFAULT_CFG, RANK
 
 
 class SegmentationTrainer(yolo.detect.DetectionTrainer):
-    """
-    A class extending the DetectionTrainer class for training based on a segmentation model.
+    """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.
@@ -28,8 +27,7 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks=None):
-        """
-        Initialize a SegmentationTrainer object.
+        """Initialize a SegmentationTrainer object.
 
         Args:
             cfg (dict): Configuration dictionary with default training settings.
@@ -42,8 +40,7 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer):
         super().__init__(cfg, overrides, _callbacks)
 
     def get_model(self, cfg: dict | str | None = None, weights: str | Path | None = None, verbose: bool = True):
-        """
-        Initialize and return a SegmentationModel with specified configuration and weights.
+        """Initialize and return a SegmentationModel with specified configuration and weights.
 
         Args:
             cfg (dict | str, optional): Model configuration. Can be a dictionary, a path to a YAML file, or None.

ultralytics/models/yolo/segment/val.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from multiprocessing.pool import ThreadPool
 from pathlib import Path
 from typing import Any
 
@@ -11,17 +10,16 @@ import torch
 import torch.nn.functional as F
 
 from ultralytics.models.yolo.detect import DetectionValidator
-from ultralytics.utils import LOGGER, NUM_THREADS, ops
+from ultralytics.utils import LOGGER, ops
 from ultralytics.utils.checks import check_requirements
 from ultralytics.utils.metrics import SegmentMetrics, mask_iou
 
 
 class SegmentationValidator(DetectionValidator):
-    """
-    A class extending the DetectionValidator class for validation based on a segmentation model.
+    """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.
+    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.
@@ -38,11 +36,10 @@ class SegmentationValidator(DetectionValidator):
     """
 
     def __init__(self, dataloader=None, save_dir=None, args=None, _callbacks=None) -> 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.
+            dataloader (torch.utils.data.DataLoader, optional): DataLoader to use for validation.
             save_dir (Path, optional): Directory to save results.
             args (namespace, optional): Arguments for the validator.
             _callbacks (list, optional): List of callback functions.
@@ -53,8 +50,7 @@ class SegmentationValidator(DetectionValidator):
         self.metrics = SegmentMetrics()
 
     def preprocess(self, batch: dict[str, Any]) -> dict[str, Any]:
-        """
-        Preprocess batch of images for YOLO segmentation validation.
+        """Preprocess batch of images for YOLO segmentation validation.
 
         Args:
             batch (dict[str, Any]): Batch containing images and annotations.
@@ -67,8 +63,7 @@ class SegmentationValidator(DetectionValidator):
         return batch
 
     def init_metrics(self, model: torch.nn.Module) -> None:
-        """
-        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.
@@ -96,8 +91,7 @@ class SegmentationValidator(DetectionValidator):
         )
 
     def postprocess(self, preds: list[torch.Tensor]) -> list[dict[str, torch.Tensor]]:
-        """
-        Post-process YOLO predictions and return output detections with proto.
+        """Post-process YOLO predictions and return output detections with proto.
 
         Args:
             preds (list[torch.Tensor]): Raw predictions from the model.
@@ -122,8 +116,7 @@ class SegmentationValidator(DetectionValidator):
         return preds
 
     def _prepare_batch(self, si: int, batch: dict[str, Any]) -> dict[str, Any]:
-        """
-        Prepare 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.
@@ -149,8 +142,7 @@ class SegmentationValidator(DetectionValidator):
         return prepared_batch
 
     def _process_batch(self, preds: dict[str, torch.Tensor], batch: dict[str, Any]) -> dict[str, np.ndarray]:
-        """
-        Compute correct prediction matrix for a batch based on bounding boxes and optional masks.
+        """Compute correct prediction matrix for a batch based on bounding boxes and optional masks.
 
         Args:
             preds (dict[str, torch.Tensor]): Dictionary containing predictions with keys like 'cls' and 'masks'.
@@ -159,28 +151,27 @@ class SegmentationValidator(DetectionValidator):
         Returns:
             (dict[str, np.ndarray]): A dictionary containing correct prediction matrices including 'tp_m' for mask IoU.
 
-        Notes:
-            - 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.
-
         Examples:
             >>> preds = {"cls": torch.tensor([1, 0]), "masks": torch.rand(2, 640, 640), "bboxes": torch.rand(2, 4)}
             >>> batch = {"cls": torch.tensor([1, 0]), "masks": torch.rand(2, 640, 640), "bboxes": torch.rand(2, 4)}
             >>> correct_preds = validator._process_batch(preds, batch)
+
+        Notes:
+            - 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.
         """
         tp = super()._process_batch(preds, batch)
         gt_cls = batch["cls"]
         if gt_cls.shape[0] == 0 or preds["cls"].shape[0] == 0:
             tp_m = np.zeros((preds["cls"].shape[0], self.niou), dtype=bool)
         else:
-            iou = mask_iou(batch["masks"].flatten(1), preds["masks"].flatten(1))
+            iou = mask_iou(batch["masks"].flatten(1), preds["masks"].flatten(1).float())  # float, uint8
             tp_m = self.match_predictions(preds["cls"], gt_cls, iou).cpu().numpy()
         tp.update({"tp_m": tp_m})  # update tp with mask IoU
         return tp
 
     def plot_predictions(self, batch: dict[str, Any], preds: list[dict[str, torch.Tensor]], ni: int) -> None:
-        """
-        Plot batch predictions with masks and bounding boxes.
+        """Plot batch predictions with masks and bounding boxes.
 
         Args:
             batch (dict[str, Any]): Batch containing images and annotations.
@@ -195,8 +186,7 @@ class SegmentationValidator(DetectionValidator):
         super().plot_predictions(batch, preds, ni, max_det=self.args.max_det)  # plot bboxes
 
     def save_one_txt(self, predn: torch.Tensor, save_conf: bool, shape: tuple[int, int], file: Path) -> None:
-        """
-        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, class).
@@ -215,24 +205,84 @@ class SegmentationValidator(DetectionValidator):
         ).save_txt(file, save_conf=save_conf)
 
     def pred_to_json(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> None:
-        """
-        Save one JSON result for COCO evaluation.
+        """Save one JSON result for COCO evaluation.
 
         Args:
             predn (dict[str, torch.Tensor]): Predictions containing bboxes, masks, confidence scores, and classes.
             pbatch (dict[str, Any]): Batch dictionary containing 'imgsz', 'ori_shape', 'ratio_pad', and 'im_file'.
         """
-        from faster_coco_eval.core.mask import encode  # noqa
-
-        def single_encode(x):
-            """Encode predicted masks as RLE and append results to jdict."""
-            rle = encode(np.asarray(x[:, :, None], order="F", dtype="uint8"))[0]
-            rle["counts"] = rle["counts"].decode("utf-8")
-            return rle
 
-        pred_masks = np.transpose(predn["masks"], (2, 0, 1))
-        with ThreadPool(NUM_THREADS) as pool:
-            rles = pool.map(single_encode, pred_masks)
+        def to_string(counts: list[int]) -> str:
+            """Converts the RLE object into a compact string representation. Each count is delta-encoded and
+            variable-length encoded as a string.
+
+            Args:
+                counts (list[int]): List of RLE counts.
+            """
+            result = []
+
+            for i in range(len(counts)):
+                x = int(counts[i])
+
+                # Apply delta encoding for all counts after the second entry
+                if i > 2:
+                    x -= int(counts[i - 2])
+
+                # Variable-length encode the value
+                while True:
+                    c = x & 0x1F  # Take 5 bits
+                    x >>= 5
+
+                    # If the sign bit (0x10) is set, continue if x != -1;
+                    # otherwise, continue if x != 0
+                    more = (x != -1) if (c & 0x10) else (x != 0)
+                    if more:
+                        c |= 0x20  # Set continuation bit
+                    c += 48  # Shift to ASCII
+                    result.append(chr(c))
+                    if not more:
+                        break
+
+            return "".join(result)
+
+        def multi_encode(pixels: torch.Tensor) -> list[int]:
+            """Convert multiple binary masks using Run-Length Encoding (RLE).
+
+            Args:
+                pixels (torch.Tensor): A 2D tensor where each row represents a flattened binary mask with shape [N,
+                    H*W].
+
+            Returns:
+                (list[int]): A list of RLE counts for each mask.
+            """
+            transitions = pixels[:, 1:] != pixels[:, :-1]
+            row_idx, col_idx = torch.where(transitions)
+            col_idx = col_idx + 1
+
+            # Compute run lengths
+            counts = []
+            for i in range(pixels.shape[0]):
+                positions = col_idx[row_idx == i]
+                if len(positions):
+                    count = torch.diff(positions).tolist()
+                    count.insert(0, positions[0].item())
+                    count.append(len(pixels[i]) - positions[-1].item())
+                else:
+                    count = [len(pixels[i])]
+
+                # Ensure starting with background (0) count
+                if pixels[i][0].item() == 1:
+                    count = [0, *count]
+                counts.append(count)
+
+            return counts
+
+        pred_masks = predn["masks"].transpose(2, 1).contiguous().view(len(predn["masks"]), -1)  # N, H*W
+        h, w = predn["masks"].shape[1:3]
+        counts = multi_encode(pred_masks)
+        rles = []
+        for c in counts:
+            rles.append({"size": [h, w], "counts": to_string(c)})
         super().pred_to_json(predn, pbatch)
         for i, r in enumerate(rles):
             self.jdict[-len(rles) + i]["segmentation"] = r  # segmentation
@@ -241,11 +291,9 @@ class SegmentationValidator(DetectionValidator):
         """Scales predictions to the original image size."""
         return {
             **super().scale_preds(predn, pbatch),
-            "masks": ops.scale_image(
-                torch.as_tensor(predn["masks"], dtype=torch.uint8).permute(1, 2, 0).contiguous().cpu().numpy(),
-                pbatch["ori_shape"],
-                ratio_pad=pbatch["ratio_pad"],
-            ),
+            "masks": ops.scale_masks(predn["masks"][None], pbatch["ori_shape"], ratio_pad=pbatch["ratio_pad"])[
+                0
+            ].byte(),
         }
 
     def eval_json(self, stats: dict[str, Any]) -> dict[str, Any]:

ultralytics/models/yolo/world/train.py

@@ -24,8 +24,7 @@ def on_pretrain_routine_end(trainer) -> None:
 
 
 class WorldTrainer(DetectionTrainer):
-    """
-    A trainer class for fine-tuning YOLO World models on close-set datasets.
+    """A trainer class for fine-tuning YOLO World models on close-set datasets.
 
     This trainer extends the DetectionTrainer to support training YOLO World models, which combine visual and textual
     features for improved object detection and understanding. It handles text embedding generation and caching to
@@ -54,8 +53,7 @@ class WorldTrainer(DetectionTrainer):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides: dict[str, Any] | None = None, _callbacks=None):
-        """
-        Initialize a WorldTrainer object with given arguments.
+        """Initialize a WorldTrainer object with given arguments.
 
         Args:
             cfg (dict[str, Any]): Configuration for the trainer.
@@ -69,8 +67,7 @@ class WorldTrainer(DetectionTrainer):
         self.text_embeddings = None
 
     def get_model(self, cfg=None, weights: str | None = None, verbose: bool = True) -> WorldModel:
-        """
-        Return WorldModel initialized with specified config and weights.
+        """Return WorldModel initialized with specified config and weights.
 
         Args:
             cfg (dict[str, Any] | str, optional): Model configuration.
@@ -95,8 +92,7 @@ class WorldTrainer(DetectionTrainer):
         return model
 
     def build_dataset(self, img_path: str, mode: str = "train", batch: int | None = None):
-        """
-        Build YOLO Dataset for training or validation.
+        """Build YOLO Dataset for training or validation.
 
         Args:
             img_path (str): Path to the folder containing images.
@@ -115,11 +111,10 @@ class WorldTrainer(DetectionTrainer):
         return dataset
 
     def set_text_embeddings(self, datasets: list[Any], batch: int | None) -> None:
-        """
-        Set text embeddings for datasets to accelerate training by caching category names.
+        """Set text embeddings for datasets to accelerate training by caching category names.
 
-        This method collects unique category names from all datasets, then generates and caches text embeddings
-        for these categories to improve training efficiency.
+        This method collects unique category names from all datasets, then generates and caches text embeddings for
+        these categories to improve training efficiency.
 
         Args:
             datasets (list[Any]): List of datasets from which to extract category names.
@@ -141,8 +136,7 @@ class WorldTrainer(DetectionTrainer):
         self.text_embeddings = text_embeddings
 
     def generate_text_embeddings(self, texts: list[str], batch: int, cache_dir: Path) -> dict[str, torch.Tensor]:
-        """
-        Generate text embeddings for a list of text samples.
+        """Generate text embeddings for a list of text samples.
 
         Args:
             texts (list[str]): List of text samples to encode.

ultralytics/models/yolo/world/train_world.py

@@ -10,8 +10,7 @@ from ultralytics.utils.torch_utils import unwrap_model
 
 
 class WorldTrainerFromScratch(WorldTrainer):
-    """
-    A class extending the WorldTrainer for training a world model from scratch on open-set datasets.
+    """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.
@@ -53,45 +52,25 @@ class WorldTrainerFromScratch(WorldTrainer):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """
-        Initialize a WorldTrainerFromScratch object.
+        """Initialize a WorldTrainerFromScratch object.
 
-        This initializes a trainer for YOLO-World models from scratch, supporting mixed datasets including both
-        object detection and grounding datasets for vision-language capabilities.
+        This initializes a trainer for YOLO-World models from scratch, supporting mixed datasets including both object
+        detection and grounding datasets for vision-language capabilities.
 
         Args:
             cfg (dict): Configuration dictionary with default parameters for model training.
             overrides (dict, optional): Dictionary of parameter overrides to customize the configuration.
             _callbacks (list, optional): List of callback functions to be executed during different stages of training.
-
-        Examples:
-            >>> from ultralytics.models.yolo.world.train_world import WorldTrainerFromScratch
-            >>> from ultralytics import YOLOWorld
-            >>> data = dict(
-            ...     train=dict(
-            ...         yolo_data=["Objects365.yaml"],
-            ...         grounding_data=[
-            ...             dict(
-            ...                 img_path="flickr30k/images",
-            ...                 json_file="flickr30k/final_flickr_separateGT_train.json",
-            ...             ),
-            ...         ],
-            ...     ),
-            ...     val=dict(yolo_data=["lvis.yaml"]),
-            ... )
-            >>> model = YOLOWorld("yolov8s-worldv2.yaml")
-            >>> model.train(data=data, trainer=WorldTrainerFromScratch)
         """
         if overrides is None:
             overrides = {}
         super().__init__(cfg, overrides, _callbacks)
 
     def build_dataset(self, img_path, mode="train", batch=None):
-        """
-        Build YOLO Dataset for training or validation.
+        """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.
+        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 or list of paths.
@@ -122,11 +101,10 @@ class WorldTrainerFromScratch(WorldTrainer):
         return YOLOConcatDataset(datasets) if len(datasets) > 1 else datasets[0]
 
     def get_dataset(self):
-        """
-        Get train and validation paths from data dictionary.
+        """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.
+        Processes the data configuration to extract paths for training and validation datasets, handling both YOLO
+        detection datasets and grounding datasets.
 
         Returns:
             train_path (str): Train dataset path.
@@ -187,8 +165,7 @@ class WorldTrainerFromScratch(WorldTrainer):
         pass
 
     def final_eval(self):
-        """
-        Perform final evaluation and validation for the 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.
 

ultralytics/models/yolo/yoloe/__init__.py

@@ -6,17 +6,17 @@ from .train_seg import YOLOEPESegTrainer, YOLOESegTrainer, YOLOESegTrainerFromSc
 from .val import YOLOEDetectValidator, YOLOESegValidator
 
 __all__ = [
-    "YOLOETrainer",
-    "YOLOEPETrainer",
-    "YOLOESegTrainer",
     "YOLOEDetectValidator",
-    "YOLOESegValidator",
+    "YOLOEPEFreeTrainer",
     "YOLOEPESegTrainer",
+    "YOLOEPETrainer",
+    "YOLOESegTrainer",
     "YOLOESegTrainerFromScratch",
     "YOLOESegVPTrainer",
-    "YOLOEVPTrainer",
-    "YOLOEPEFreeTrainer",
+    "YOLOESegValidator",
+    "YOLOETrainer",
+    "YOLOETrainerFromScratch",
     "YOLOEVPDetectPredictor",
     "YOLOEVPSegPredictor",
-    "YOLOETrainerFromScratch",
+    "YOLOEVPTrainer",
 ]

ultralytics/models/yolo/yoloe/predict.py

@@ -9,11 +9,10 @@ from ultralytics.models.yolo.segment import SegmentationPredictor
 
 
 class YOLOEVPDetectPredictor(DetectionPredictor):
-    """
-    A mixin class for YOLO-EVP (Enhanced Visual Prompting) predictors.
+    """A mixin class for YOLO-EVP (Enhanced Visual Prompting) predictors.
 
-    This mixin provides common functionality for YOLO models that use visual prompting, including
-    model setup, prompt handling, and preprocessing transformations.
+    This mixin provides common functionality for YOLO models that use visual prompting, including model setup, prompt
+    handling, and preprocessing transformations.
 
     Attributes:
         model (torch.nn.Module): The YOLO model for inference.
@@ -29,8 +28,7 @@ class YOLOEVPDetectPredictor(DetectionPredictor):
     """
 
     def setup_model(self, model, verbose: bool = True):
-        """
-        Set up the model for prediction.
+        """Set up the model for prediction.
 
         Args:
             model (torch.nn.Module): Model to load or use.
@@ -40,21 +38,19 @@ class YOLOEVPDetectPredictor(DetectionPredictor):
         self.done_warmup = True
 
     def set_prompts(self, prompts):
-        """
-        Set the visual prompts for the model.
+        """Set the visual prompts for the model.
 
         Args:
-            prompts (dict): Dictionary containing class indices and bounding boxes or masks.
-                Must include a 'cls' key with class indices.
+            prompts (dict): Dictionary containing class indices and bounding boxes or masks. Must include a 'cls' key
+                with class indices.
         """
         self.prompts = prompts
 
     def pre_transform(self, im):
-        """
-        Preprocess images and prompts before inference.
+        """Preprocess images and prompts before inference.
 
-        This method applies letterboxing to the input image and transforms the visual prompts
-        (bounding boxes or masks) accordingly.
+        This method applies letterboxing to the input image and transforms the visual prompts (bounding boxes or masks)
+        accordingly.
 
         Args:
             im (list): List containing a single input image.
@@ -94,8 +90,7 @@ class YOLOEVPDetectPredictor(DetectionPredictor):
         return img
 
     def _process_single_image(self, dst_shape, src_shape, category, bboxes=None, masks=None):
-        """
-        Process a single image by resizing bounding boxes or masks and generating visuals.
+        """Process a single image by resizing bounding boxes or masks and generating visuals.
 
         Args:
             dst_shape (tuple): The target shape (height, width) of the image.
@@ -131,8 +126,7 @@ class YOLOEVPDetectPredictor(DetectionPredictor):
         return LoadVisualPrompt().get_visuals(category, dst_shape, bboxes, masks)
 
     def inference(self, im, *args, **kwargs):
-        """
-        Run inference with visual prompts.
+        """Run inference with visual prompts.
 
         Args:
             im (torch.Tensor): Input image tensor.
@@ -145,13 +139,12 @@ class YOLOEVPDetectPredictor(DetectionPredictor):
         return super().inference(im, vpe=self.prompts, *args, **kwargs)
 
     def get_vpe(self, source):
-        """
-        Process the source to get the visual prompt embeddings (VPE).
+        """Process the source to get the visual prompt embeddings (VPE).
 
         Args:
-            source (str | Path | int | PIL.Image | np.ndarray | torch.Tensor | list | tuple): The source
-                of the image to make predictions on. Accepts various types including file paths, URLs, PIL
-                images, numpy arrays, and torch tensors.
+            source (str | Path | int | PIL.Image | np.ndarray | torch.Tensor | list | tuple): The source of the image to
+                make predictions on. Accepts various types including file paths, URLs, PIL images, numpy arrays, and
+                torch tensors.
 
         Returns:
             (torch.Tensor): The visual prompt embeddings (VPE) from the model.