dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/models/yolo/model.py

@@ -1,6 +1,11 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from pathlib import Path
+from typing import Any
+
+import torch
 
 from ultralytics.data.build import load_inference_source
 from ultralytics.engine.model import Model
@@ -19,19 +24,42 @@ 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):
-        """
-        Initialize a YOLO model.
+    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")
+    """
 
-        This constructor initializes a YOLO model, automatically switching to specialized model types
-        (YOLOWorld or YOLOE) based on the model filename.
+    def __init__(self, model: str | Path = "yolo11n.pt", task: str | None = None, verbose: bool = False):
+        """Initialize a YOLO model.
+
+        This constructor initializes a YOLO model, automatically switching to specialized model types (YOLOWorld or
+        YOLOE) based on the model filename.
 
         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'.
-                Defaults to auto-detection based on model.
+            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.
 
         Examples:
@@ -39,7 +67,7 @@ class YOLO(Model):
             >>> model = YOLO("yolo11n.pt")  # load a pretrained YOLOv11n detection model
             >>> model = YOLO("yolo11n-seg.pt")  # load a pretrained YOLO11n segmentation model
         """
-        path = Path(model)
+        path = Path(model if isinstance(model, (str, Path)) else "")
         if "-world" in path.stem and path.suffix in {".pt", ".yaml", ".yml"}:  # if YOLOWorld PyTorch model
             new_instance = YOLOWorld(path, verbose=verbose)
             self.__class__ = type(new_instance)
@@ -51,9 +79,15 @@ class YOLO(Model):
         else:
             # Continue with default YOLO initialization
             super().__init__(model=model, task=task, verbose=verbose)
+            if hasattr(self.model, "model") and "RTDETR" in self.model.model[-1]._get_name():  # if RTDETR head
+                from ultralytics import RTDETR
+
+                new_instance = RTDETR(self)
+                self.__class__ = type(new_instance)
+                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": {
@@ -90,14 +124,35 @@ 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:
-        """
-        Initialize YOLOv8-World model with a pre-trained model file.
+    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: str | Path = "yolov8s-world.pt", verbose: bool = False) -> None:
+        """Initialize YOLOv8-World model with a pre-trained model file.
 
-        Loads a YOLOv8-World model for object detection. If no custom class names are provided, it assigns default
-        COCO class names.
+        Loads a YOLOv8-World model for object detection. If no custom class names are provided, it assigns default COCO
+        class names.
 
         Args:
             model (str | Path): Path to the pre-trained model file. Supports *.pt and *.yaml formats.
@@ -110,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": {
@@ -121,9 +176,8 @@ class YOLOWorld(Model):
             }
         }
 
-    def set_classes(self, classes):
-        """
-        Set the model's class names for detection.
+    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"].
@@ -141,11 +195,41 @@ 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:
-        """
-        Initialize YOLOE model with a pre-trained model file.
+    """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: str | Path = "yoloe-11s-seg.pt", task: str | None = None, verbose: bool = False) -> None:
+        """Initialize YOLOE model with a pre-trained model file.
 
         Args:
             model (str | Path): Path to the pre-trained model file. Supports *.pt and *.yaml formats.
@@ -154,12 +238,8 @@ class YOLOE(Model):
         """
         super().__init__(model=model, task=task, verbose=verbose)
 
-        # Assign default COCO class names when there are no custom names
-        if not hasattr(self.model, "names"):
-            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": {
@@ -182,11 +262,10 @@ class YOLOE(Model):
         return self.model.get_text_pe(texts)
 
     def get_visual_pe(self, img, visual):
-        """
-        Get visual positional embeddings for the given image and visual features.
+        """Get visual positional embeddings for the given image and visual features.
 
-        This method extracts positional embeddings from visual features based on the input image. It requires
-        that the model is an instance of YOLOEModel.
+        This method extracts positional embeddings from visual features based on the input image. It requires that the
+        model is an instance of YOLOEModel.
 
         Args:
             img (torch.Tensor): Input image tensor.
@@ -198,22 +277,21 @@ class YOLOE(Model):
         Examples:
             >>> model = YOLOE("yoloe-11s-seg.pt")
             >>> img = torch.rand(1, 3, 640, 640)
-            >>> visual_features = model.model.backbone(img)
+            >>> visual_features = torch.rand(1, 1, 80, 80)
             >>> pe = model.get_visual_pe(img, visual_features)
         """
         assert isinstance(self.model, YOLOEModel)
         return self.model.get_visual_pe(img, visual)
 
-    def set_vocab(self, vocab, names):
-        """
-        Set vocabulary and class names for the YOLOE model.
+    def set_vocab(self, vocab: list[str], names: list[str]) -> None:
+        """Set vocabulary and class names for the YOLOE model.
 
-        This method configures the vocabulary and class names used by the model for text processing and
-        classification tasks. The model must be an instance of YOLOEModel.
+        This method configures the vocabulary and class names used by the model for text processing and 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.
@@ -230,15 +308,16 @@ class YOLOE(Model):
         assert isinstance(self.model, YOLOEModel)
         return self.model.get_vocab(names)
 
-    def set_classes(self, classes, embeddings):
-        """
-        Set the model's class names and embeddings for detection.
+    def set_classes(self, classes: list[str], embeddings: torch.Tensor | None = None) -> None:
+        """Set the model's class names and embeddings for detection.
 
         Args:
             classes (list[str]): A list of categories i.e. ["person"].
             embeddings (torch.Tensor): Embeddings corresponding to the classes.
         """
         assert isinstance(self.model, YOLOEModel)
+        if embeddings is None:
+            embeddings = self.get_text_pe(classes)  # generate text embeddings if not provided
         self.model.set_classes(classes, embeddings)
         # Verify no background class is present
         assert " " not in classes
@@ -251,12 +330,11 @@ class YOLOE(Model):
     def val(
         self,
         validator=None,
-        load_vp=False,
-        refer_data=None,
+        load_vp: bool = False,
+        refer_data: str | None = None,
         **kwargs,
     ):
-        """
-        Validate the model using text or visual prompts.
+        """Validate the model using text or visual prompts.
 
         Args:
             validator (callable, optional): A callable validator function. If None, a default validator is loaded.
@@ -279,28 +357,27 @@ class YOLOE(Model):
         self,
         source=None,
         stream: bool = False,
-        visual_prompts: dict = {},
+        visual_prompts: dict[str, list] = {},
         refer_image=None,
-        predictor=None,
+        predictor=yolo.yoloe.YOLOEVPDetectPredictor,
         **kwargs,
     ):
-        """
-        Run prediction on images, videos, directories, streams, etc.
+        """Run prediction on images, videos, directories, streams, etc.
 
         Args:
-            source (str | int | PIL.Image | np.ndarray, optional): Source for prediction. Accepts image paths,
-                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.
+            source (str | int | PIL.Image | np.ndarray, optional): Source for prediction. Accepts image paths, 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[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.
+            predictor (callable, optional): Custom predictor function. If None, a predictor is automatically loaded
+                based on the task.
             **kwargs (Any): Additional keyword arguments passed to the predictor.
 
         Returns:
-            (List | generator): List of Results objects or generator of Results objects if stream=True.
+            (list | generator): List of Results objects or generator of Results objects if stream=True.
 
         Examples:
             >>> model = YOLOE("yoloe-11s-seg.pt")
@@ -317,18 +394,21 @@ class YOLOE(Model):
                 f"Expected equal number of bounding boxes and classes, but got {len(visual_prompts['bboxes'])} and "
                 f"{len(visual_prompts['cls'])} respectively"
             )
-        self.predictor = (predictor or self._smart_load("predictor"))(
-            overrides={
-                "task": self.model.task,
-                "mode": "predict",
-                "save": False,
-                "verbose": refer_image is None,
-                "batch": 1,
-            },
-            _callbacks=self.callbacks,
-        )
+            if type(self.predictor) is not predictor:
+                self.predictor = predictor(
+                    overrides={
+                        "task": self.model.task,
+                        "mode": "predict",
+                        "save": False,
+                        "verbose": refer_image is None,
+                        "batch": 1,
+                        "device": kwargs.get("device", None),
+                        "half": kwargs.get("half", False),
+                        "imgsz": kwargs.get("imgsz", self.overrides["imgsz"]),
+                    },
+                    _callbacks=self.callbacks,
+                )
 
-        if len(visual_prompts):
             num_cls = (
                 max(len(set(c)) for c in visual_prompts["cls"])
                 if isinstance(source, list) and refer_image is None  # means multiple images
@@ -337,18 +417,19 @@ class YOLOE(Model):
             self.model.model[-1].nc = num_cls
             self.model.names = [f"object{i}" for i in range(num_cls)]
             self.predictor.set_prompts(visual_prompts.copy())
-
-        self.predictor.setup_model(model=self.model)
-
-        if refer_image is None and source is not None:
-            dataset = load_inference_source(source)
-            if dataset.mode in {"video", "stream"}:
-                # NOTE: set the first frame as refer image for videos/streams inference
-                refer_image = next(iter(dataset))[1][0]
-        if refer_image is not None and len(visual_prompts):
-            vpe = self.predictor.get_vpe(refer_image)
-            self.model.set_classes(self.model.names, vpe)
-            self.task = "segment" if isinstance(self.predictor, yolo.segment.SegmentationPredictor) else "detect"
-            self.predictor = None  # reset predictor
+            self.predictor.setup_model(model=self.model)
+
+            if refer_image is None and source is not None:
+                dataset = load_inference_source(source)
+                if dataset.mode in {"video", "stream"}:
+                    # NOTE: set the first frame as refer image for videos/streams inference
+                    refer_image = next(iter(dataset))[1][0]
+            if refer_image is not None:
+                vpe = self.predictor.get_vpe(refer_image)
+                self.model.set_classes(self.model.names, vpe)
+                self.task = "segment" if isinstance(self.predictor, yolo.segment.SegmentationPredictor) else "detect"
+                self.predictor = None  # reset predictor
+        elif isinstance(self.predictor, yolo.yoloe.YOLOEVPDetectPredictor):
+            self.predictor = None  # reset predictor if no visual prompts
 
         return super().predict(source, stream, **kwargs)

ultralytics/models/yolo/obb/predict.py

@@ -8,8 +8,7 @@ from ultralytics.utils import DEFAULT_CFG, ops
 
 
 class OBBPredictor(DetectionPredictor):
-    """
-    A class extending the DetectionPredictor class for prediction based on an Oriented Bounding Box (OBB) model.
+    """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.
@@ -27,10 +26,7 @@ class OBBPredictor(DetectionPredictor):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """
-        Initialize OBBPredictor with optional model and data configuration overrides.
-
-        This constructor sets up an OBBPredictor instance for oriented bounding box detection tasks.
+        """Initialize OBBPredictor with optional model and data configuration overrides.
 
         Args:
             cfg (dict, optional): Default configuration for the predictor.
@@ -47,18 +43,18 @@ class OBBPredictor(DetectionPredictor):
         self.args.task = "obb"
 
     def construct_result(self, pred, img, orig_img, img_path):
-        """
-        Construct 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 with shape (N, 6) where
-                the last dimension contains [x, y, w, h, confidence, class_id, angle].
+            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,10 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from copy import copy
+from pathlib import Path
+from typing import Any
 
 from ultralytics.models import yolo
 from ultralytics.nn.tasks import OBBModel
@@ -8,11 +12,14 @@ from ultralytics.utils import DEFAULT_CFG, RANK
 
 
 class OBBTrainer(yolo.detect.DetectionTrainer):
-    """
-    A class extending the DetectionTrainer class for training based on an Oriented Bounding Box (OBB) model.
+    """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,39 +32,30 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
         >>> trainer.train()
     """
 
-    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
-        """
-        Initialize an OBBTrainer object for training Oriented Bounding Box (OBB) models.
-
-        This trainer extends the DetectionTrainer class to specialize in training models that detect oriented
-        bounding boxes. It automatically sets the task to 'obb' in the configuration.
+    def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks: list[Any] | None = None):
+        """Initialize an OBBTrainer object for training Oriented Bounding Box (OBB) models.
 
         Args:
-            cfg (dict, optional): Configuration dictionary for the trainer. Contains training parameters and
-                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.
-
-        Examples:
-            >>> from ultralytics.models.yolo.obb import OBBTrainer
-            >>> args = dict(model="yolo11n-obb.pt", data="dota8.yaml", epochs=3)
-            >>> trainer = OBBTrainer(overrides=args)
-            >>> trainer.train()
+            cfg (dict, optional): Configuration dictionary for the trainer. Contains training parameters and model
+                configuration.
+            overrides (dict, optional): Dictionary of parameter overrides for the configuration. Any values here will
+                take precedence over those in cfg.
+            _callbacks (list[Any], optional): List of callback functions to be invoked during training.
         """
         if overrides is None:
             overrides = {}
         overrides["task"] = "obb"
         super().__init__(cfg, overrides, _callbacks)
 
-    def get_model(self, cfg=None, weights=None, verbose=True):
-        """
-        Return OBBModel initialized with specified config and weights.
+    def get_model(
+        self, cfg: str | dict | None = None, weights: str | Path | None = 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: