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

ultralytics/models/yolo/yoloe/train.py

@@ -19,11 +19,10 @@ from .val import YOLOEDetectValidator
 
 
 class YOLOETrainer(DetectionTrainer):
-    """
-    A trainer class for YOLOE object detection models.
+    """A trainer class for YOLOE object detection models.
 
-    This class extends DetectionTrainer to provide specialized training functionality for YOLOE models,
-    including custom model initialization, validation, and dataset building with multi-modal support.
+    This class extends DetectionTrainer to provide specialized training functionality for YOLOE models, including custom
+    model initialization, validation, and dataset building with multi-modal support.
 
     Attributes:
         loss_names (tuple): Names of loss components used during training.
@@ -35,8 +34,7 @@ class YOLOETrainer(DetectionTrainer):
     """
 
     def __init__(self, cfg=DEFAULT_CFG, overrides: dict | None = None, _callbacks=None):
-        """
-        Initialize the YOLOE Trainer with specified configurations.
+        """Initialize the YOLOE Trainer with specified configurations.
 
         Args:
             cfg (dict): Configuration dictionary with default training settings from DEFAULT_CFG.
@@ -50,12 +48,11 @@ class YOLOETrainer(DetectionTrainer):
         super().__init__(cfg, overrides, _callbacks)
 
     def get_model(self, cfg=None, weights=None, verbose: bool = True):
-        """
-        Return a YOLOEModel initialized with the specified configuration and weights.
+        """Return a YOLOEModel initialized with the specified configuration and weights.
 
         Args:
-            cfg (dict | str, optional): Model configuration. Can be a dictionary containing a 'yaml_file' key,
-                a direct path to a YAML file, or None to use default configuration.
+            cfg (dict | str, optional): Model configuration. Can be a dictionary containing a 'yaml_file' key, a direct
+                path to a YAML file, or None to use default configuration.
             weights (str | Path, optional): Path to pretrained weights file to load into the model.
             verbose (bool): Whether to display model information during initialization.
 
@@ -88,8 +85,7 @@ class YOLOETrainer(DetectionTrainer):
         )
 
     def build_dataset(self, img_path: str, mode: str = "train", batch: int | None = None):
-        """
-        Build YOLO Dataset.
+        """Build YOLO Dataset.
 
         Args:
             img_path (str): Path to the folder containing images.
@@ -106,19 +102,17 @@ class YOLOETrainer(DetectionTrainer):
 
 
 class YOLOEPETrainer(DetectionTrainer):
-    """
-    Fine-tune YOLOE model using linear probing approach.
+    """Fine-tune YOLOE model using linear probing approach.
 
-    This trainer freezes most model layers and only trains specific projection layers for efficient
-    fine-tuning on new datasets while preserving pretrained features.
+    This trainer freezes most model layers and only trains specific projection layers for efficient fine-tuning on new
+    datasets while preserving pretrained features.
 
     Methods:
         get_model: Initialize YOLOEModel with frozen layers except projection layers.
     """
 
     def get_model(self, cfg=None, weights=None, verbose: bool = True):
-        """
-        Return YOLOEModel initialized with specified config and weights.
+        """Return YOLOEModel initialized with specified config and weights.
 
         Args:
             cfg (dict | str, optional): Model configuration.
@@ -153,18 +147,22 @@ class YOLOEPETrainer(DetectionTrainer):
         model.model[-1].cv3[0][2] = deepcopy(model.model[-1].cv3[0][2]).requires_grad_(True)
         model.model[-1].cv3[1][2] = deepcopy(model.model[-1].cv3[1][2]).requires_grad_(True)
         model.model[-1].cv3[2][2] = deepcopy(model.model[-1].cv3[2][2]).requires_grad_(True)
-        del model.pe
+
+        if getattr(model.model[-1], "one2one_cv3", None) is not None:
+            model.model[-1].one2one_cv3[0][2] = deepcopy(model.model[-1].cv3[0][2]).requires_grad_(True)
+            model.model[-1].one2one_cv3[1][2] = deepcopy(model.model[-1].cv3[1][2]).requires_grad_(True)
+            model.model[-1].one2one_cv3[2][2] = deepcopy(model.model[-1].cv3[2][2]).requires_grad_(True)
+
         model.train()
 
         return model
 
 
 class YOLOETrainerFromScratch(YOLOETrainer, WorldTrainerFromScratch):
-    """
-    Train YOLOE models from scratch with text embedding support.
+    """Train YOLOE models from scratch with text embedding support.
 
-    This trainer combines YOLOE training capabilities with world training features, enabling
-    training from scratch with text embeddings and grounding datasets.
+    This trainer combines YOLOE training capabilities with world training features, enabling training from scratch with
+    text embeddings and grounding datasets.
 
     Methods:
         build_dataset: Build datasets for training with grounding support.
@@ -172,11 +170,10 @@ class YOLOETrainerFromScratch(YOLOETrainer, WorldTrainerFromScratch):
     """
 
     def build_dataset(self, img_path: list[str] | str, mode: str = "train", batch: int | None = 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.
@@ -189,8 +186,7 @@ class YOLOETrainerFromScratch(YOLOETrainer, WorldTrainerFromScratch):
         return WorldTrainerFromScratch.build_dataset(self, img_path, mode, batch)
 
     def generate_text_embeddings(self, texts: list[str], batch: int, cache_dir: Path):
-        """
-        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.
@@ -216,11 +212,10 @@ class YOLOETrainerFromScratch(YOLOETrainer, WorldTrainerFromScratch):
 
 
 class YOLOEPEFreeTrainer(YOLOEPETrainer, YOLOETrainerFromScratch):
-    """
-    Train prompt-free YOLOE model.
+    """Train prompt-free YOLOE model.
 
-    This trainer combines linear probing capabilities with from-scratch training for prompt-free
-    YOLOE models that don't require text prompts during inference.
+    This trainer combines linear probing capabilities with from-scratch training for prompt-free YOLOE models that don't
+    require text prompts during inference.
 
     Methods:
         get_validator: Return standard DetectionValidator for validation.
@@ -240,12 +235,11 @@ class YOLOEPEFreeTrainer(YOLOEPETrainer, YOLOETrainerFromScratch):
         return DetectionTrainer.preprocess_batch(self, batch)
 
     def set_text_embeddings(self, datasets, batch: int):
-        """
-        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, generates text embeddings for them,
-        and caches these embeddings to improve training efficiency. The embeddings are stored in a file
-        in the parent directory of the first dataset's image path.
+        This method collects unique category names from all datasets, generates text embeddings for them, and caches
+        these embeddings to improve training efficiency. The embeddings are stored in a file in the parent directory of
+        the first dataset's image path.
 
         Args:
             datasets (list[Dataset]): List of datasets containing category names to process.
@@ -260,19 +254,17 @@ class YOLOEPEFreeTrainer(YOLOEPETrainer, YOLOETrainerFromScratch):
 
 
 class YOLOEVPTrainer(YOLOETrainerFromScratch):
-    """
-    Train YOLOE model with visual prompts.
+    """Train YOLOE model with visual prompts.
 
-    This trainer extends YOLOETrainerFromScratch to support visual prompt-based training,
-    where visual cues are provided alongside images to guide the detection process.
+    This trainer extends YOLOETrainerFromScratch to support visual prompt-based training, where visual cues are provided
+    alongside images to guide the detection process.
 
     Methods:
         build_dataset: Build dataset with visual prompt loading transforms.
     """
 
     def build_dataset(self, img_path: list[str] | str, mode: str = "train", batch: int | None = None):
-        """
-        Build YOLO Dataset for training or validation with visual prompts.
+        """Build YOLO Dataset for training or validation with visual prompts.
 
         Args:
             img_path (list[str] | str): Path to the folder containing images or list of paths.

ultralytics/models/yolo/yoloe/train_seg.py

@@ -11,8 +11,7 @@ from .val import YOLOESegValidator
 
 
 class YOLOESegTrainer(YOLOETrainer, SegmentationTrainer):
-    """
-    Trainer class for YOLOE segmentation models.
+    """Trainer class for YOLOE segmentation models.
 
     This class combines YOLOETrainer and SegmentationTrainer to provide training functionality specifically for YOLOE
     segmentation models, enabling both object detection and instance segmentation capabilities.
@@ -24,8 +23,7 @@ class YOLOESegTrainer(YOLOETrainer, SegmentationTrainer):
     """
 
     def get_model(self, cfg=None, weights=None, verbose=True):
-        """
-        Return YOLOESegModel initialized with specified config and weights.
+        """Return YOLOESegModel initialized with specified config and weights.
 
         Args:
             cfg (dict | str, optional): Model configuration dictionary or YAML file path.
@@ -49,8 +47,7 @@ class YOLOESegTrainer(YOLOETrainer, SegmentationTrainer):
         return model
 
     def get_validator(self):
-        """
-        Create and return a validator for YOLOE segmentation model evaluation.
+        """Create and return a validator for YOLOE segmentation model evaluation.
 
         Returns:
             (YOLOESegValidator): Validator for YOLOE segmentation models.
@@ -62,8 +59,7 @@ class YOLOESegTrainer(YOLOETrainer, SegmentationTrainer):
 
 
 class YOLOEPESegTrainer(SegmentationTrainer):
-    """
-    Fine-tune YOLOESeg model in linear probing way.
+    """Fine-tune YOLOESeg model in linear probing way.
 
     This trainer specializes in fine-tuning YOLOESeg models using a linear probing approach, which involves freezing
     most of the model and only training specific layers for efficient adaptation to new tasks.
@@ -73,8 +69,7 @@ class YOLOEPESegTrainer(SegmentationTrainer):
     """
 
     def get_model(self, cfg=None, weights=None, verbose=True):
-        """
-        Return YOLOESegModel initialized with specified config and weights for linear probing.
+        """Return YOLOESegModel initialized with specified config and weights for linear probing.
 
         Args:
             cfg (dict | str, optional): Model configuration dictionary or YAML file path.
@@ -109,7 +104,12 @@ class YOLOEPESegTrainer(SegmentationTrainer):
         model.model[-1].cv3[0][2] = deepcopy(model.model[-1].cv3[0][2]).requires_grad_(True)
         model.model[-1].cv3[1][2] = deepcopy(model.model[-1].cv3[1][2]).requires_grad_(True)
         model.model[-1].cv3[2][2] = deepcopy(model.model[-1].cv3[2][2]).requires_grad_(True)
-        del model.pe
+
+        if getattr(model.model[-1], "one2one_cv3", None) is not None:
+            model.model[-1].one2one_cv3[0][2] = deepcopy(model.model[-1].cv3[0][2]).requires_grad_(True)
+            model.model[-1].one2one_cv3[1][2] = deepcopy(model.model[-1].cv3[1][2]).requires_grad_(True)
+            model.model[-1].one2one_cv3[2][2] = deepcopy(model.model[-1].cv3[2][2]).requires_grad_(True)
+
         model.train()
 
         return model

ultralytics/models/yolo/yoloe/val.py

@@ -21,12 +21,11 @@ from ultralytics.utils.torch_utils import select_device, smart_inference_mode
 
 
 class YOLOEDetectValidator(DetectionValidator):
-    """
-    A validator class for YOLOE detection models that handles both text and visual prompt embeddings.
+    """A validator class for YOLOE detection models that handles both text and visual prompt embeddings.
 
-    This class extends DetectionValidator to provide specialized validation functionality for YOLOE models.
-    It supports validation using either text prompts or visual prompt embeddings extracted from training samples,
-    enabling flexible evaluation strategies for prompt-based object detection.
+    This class extends DetectionValidator to provide specialized validation functionality for YOLOE models. It supports
+    validation using either text prompts or visual prompt embeddings extracted from training samples, enabling flexible
+    evaluation strategies for prompt-based object detection.
 
     Attributes:
         device (torch.device): The device on which validation is performed.
@@ -50,12 +49,11 @@ class YOLOEDetectValidator(DetectionValidator):
 
     @smart_inference_mode()
     def get_visual_pe(self, dataloader: torch.utils.data.DataLoader, model: YOLOEModel) -> torch.Tensor:
-        """
-        Extract visual prompt embeddings from training samples.
+        """Extract visual prompt embeddings from training samples.
 
-        This method processes a dataloader to compute visual prompt embeddings for each class using a YOLOE model.
-        It normalizes the embeddings and handles cases where no samples exist for a class by setting their
-        embeddings to zero.
+        This method processes a dataloader to compute visual prompt embeddings for each class using a YOLOE model. It
+        normalizes the embeddings and handles cases where no samples exist for a class by setting their embeddings to
+        zero.
 
         Args:
             dataloader (torch.utils.data.DataLoader): The dataloader providing training samples.
@@ -99,12 +97,10 @@ class YOLOEDetectValidator(DetectionValidator):
         return visual_pe.unsqueeze(0)
 
     def get_vpe_dataloader(self, data: dict[str, Any]) -> torch.utils.data.DataLoader:
-        """
-        Create a dataloader for LVIS training visual prompt samples.
+        """Create a dataloader for LVIS training visual prompt samples.
 
-        This method prepares a dataloader for visual prompt embeddings (VPE) using the specified dataset.
-        It applies necessary transformations including LoadVisualPrompt and configurations to the dataset
-        for validation purposes.
+        This method prepares a dataloader for visual prompt embeddings (VPE) using the specified dataset. It applies
+        necessary transformations including LoadVisualPrompt and configurations to the dataset for validation purposes.
 
         Args:
             data (dict): Dataset configuration dictionary containing paths and settings.
@@ -141,12 +137,11 @@ class YOLOEDetectValidator(DetectionValidator):
         refer_data: str | None = None,
         load_vp: bool = False,
     ) -> dict[str, Any]:
-        """
-        Run validation on the model using either text or visual prompt embeddings.
+        """Run validation on the model using either text or visual prompt embeddings.
 
-        This method validates the model using either text prompts or visual prompts, depending on the load_vp flag.
-        It supports validation during training (using a trainer object) or standalone validation with a provided
-        model. For visual prompts, reference data can be specified to extract embeddings from a different dataset.
+        This method validates the model using either text prompts or visual prompts, depending on the load_vp flag. It
+        supports validation during training (using a trainer object) or standalone validation with a provided model. For
+        visual prompts, reference data can be specified to extract embeddings from a different dataset.
 
         Args:
             trainer (object, optional): Trainer object containing the model and device.

ultralytics/nn/__init__.py

@@ -14,14 +14,14 @@ from .tasks import (
 )
 
 __all__ = (
+    "BaseModel",
+    "ClassificationModel",
+    "DetectionModel",
+    "SegmentationModel",
+    "guess_model_scale",
+    "guess_model_task",
     "load_checkpoint",
     "parse_model",
-    "yaml_model_load",
-    "guess_model_task",
-    "guess_model_scale",
     "torch_safe_load",
-    "DetectionModel",
-    "SegmentationModel",
-    "ClassificationModel",
-    "BaseModel",
+    "yaml_model_load",
 )