dgenerate-ultralytics-headless 8.3.190__py3-none-any.whl → 8.3.191__py3-none-any.whl

ultralytics/models/utils/loss.py

@@ -1,6 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from typing import Any, Dict, List, Optional, Tuple
+from __future__ import annotations
+
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -36,7 +38,7 @@ class DETRLoss(nn.Module):
     def __init__(
         self,
         nc: int = 80,
-        loss_gain: Optional[Dict[str, float]] = None,
+        loss_gain: dict[str, float] | None = None,
         aux_loss: bool = True,
         use_fl: bool = True,
         use_vfl: bool = False,
@@ -79,7 +81,7 @@ class DETRLoss(nn.Module):
 
     def _get_loss_class(
         self, pred_scores: torch.Tensor, targets: torch.Tensor, gt_scores: torch.Tensor, num_gts: int, postfix: str = ""
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         """
         Compute classification loss based on predictions, target values, and ground truth scores.
 
@@ -121,7 +123,7 @@ class DETRLoss(nn.Module):
 
     def _get_loss_bbox(
         self, pred_bboxes: torch.Tensor, gt_bboxes: torch.Tensor, postfix: str = ""
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         """
         Compute bounding box and GIoU losses for predicted and ground truth bounding boxes.
 
@@ -191,12 +193,12 @@ class DETRLoss(nn.Module):
         pred_scores: torch.Tensor,
         gt_bboxes: torch.Tensor,
         gt_cls: torch.Tensor,
-        gt_groups: List[int],
-        match_indices: Optional[List[Tuple]] = None,
+        gt_groups: list[int],
+        match_indices: list[tuple] | None = None,
         postfix: str = "",
-        masks: Optional[torch.Tensor] = None,
-        gt_mask: Optional[torch.Tensor] = None,
-    ) -> Dict[str, torch.Tensor]:
+        masks: torch.Tensor | None = None,
+        gt_mask: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
         """
         Get auxiliary losses for intermediate decoder layers.
 
@@ -258,7 +260,7 @@ class DETRLoss(nn.Module):
         return loss
 
     @staticmethod
-    def _get_index(match_indices: List[Tuple]) -> Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
+    def _get_index(match_indices: list[tuple]) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
         """
         Extract batch indices, source indices, and destination indices from match indices.
 
@@ -275,8 +277,8 @@ class DETRLoss(nn.Module):
         return (batch_idx, src_idx), dst_idx
 
     def _get_assigned_bboxes(
-        self, pred_bboxes: torch.Tensor, gt_bboxes: torch.Tensor, match_indices: List[Tuple]
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self, pred_bboxes: torch.Tensor, gt_bboxes: torch.Tensor, match_indices: list[tuple]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Assign predicted bounding boxes to ground truth bounding boxes based on match indices.
 
@@ -309,12 +311,12 @@ class DETRLoss(nn.Module):
         pred_scores: torch.Tensor,
         gt_bboxes: torch.Tensor,
         gt_cls: torch.Tensor,
-        gt_groups: List[int],
-        masks: Optional[torch.Tensor] = None,
-        gt_mask: Optional[torch.Tensor] = None,
+        gt_groups: list[int],
+        masks: torch.Tensor | None = None,
+        gt_mask: torch.Tensor | None = None,
         postfix: str = "",
-        match_indices: Optional[List[Tuple]] = None,
-    ) -> Dict[str, torch.Tensor]:
+        match_indices: list[tuple] | None = None,
+    ) -> dict[str, torch.Tensor]:
         """
         Calculate losses for a single prediction layer.
 
@@ -358,10 +360,10 @@ class DETRLoss(nn.Module):
         self,
         pred_bboxes: torch.Tensor,
         pred_scores: torch.Tensor,
-        batch: Dict[str, Any],
+        batch: dict[str, Any],
         postfix: str = "",
         **kwargs: Any,
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         """
         Calculate loss for predicted bounding boxes and scores.
 
@@ -407,12 +409,12 @@ class RTDETRDetectionLoss(DETRLoss):
 
     def forward(
         self,
-        preds: Tuple[torch.Tensor, torch.Tensor],
-        batch: Dict[str, Any],
-        dn_bboxes: Optional[torch.Tensor] = None,
-        dn_scores: Optional[torch.Tensor] = None,
-        dn_meta: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, torch.Tensor]:
+        preds: tuple[torch.Tensor, torch.Tensor],
+        batch: dict[str, Any],
+        dn_bboxes: torch.Tensor | None = None,
+        dn_scores: torch.Tensor | None = None,
+        dn_meta: dict[str, Any] | None = None,
+    ) -> dict[str, torch.Tensor]:
         """
         Forward pass to compute detection loss with optional denoising loss.
 
@@ -448,8 +450,8 @@ class RTDETRDetectionLoss(DETRLoss):
 
     @staticmethod
     def get_dn_match_indices(
-        dn_pos_idx: List[torch.Tensor], dn_num_group: int, gt_groups: List[int]
-    ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+        dn_pos_idx: list[torch.Tensor], dn_num_group: int, gt_groups: list[int]
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
         """
         Get match indices for denoising.
 

ultralytics/models/utils/ops.py

@@ -1,6 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from typing import Any, Dict, List, Optional, Tuple
+from __future__ import annotations
+
+from typing import Any
 
 import torch
 import torch.nn as nn
@@ -47,7 +49,7 @@ class HungarianMatcher(nn.Module):
 
     def __init__(
         self,
-        cost_gain: Optional[Dict[str, float]] = None,
+        cost_gain: dict[str, float] | None = None,
         use_fl: bool = True,
         with_mask: bool = False,
         num_sample_points: int = 12544,
@@ -82,10 +84,10 @@ class HungarianMatcher(nn.Module):
         pred_scores: torch.Tensor,
         gt_bboxes: torch.Tensor,
         gt_cls: torch.Tensor,
-        gt_groups: List[int],
-        masks: Optional[torch.Tensor] = None,
-        gt_mask: Optional[List[torch.Tensor]] = None,
-    ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+        gt_groups: list[int],
+        masks: torch.Tensor | None = None,
+        gt_mask: list[torch.Tensor] | None = None,
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
         """
         Compute optimal assignment between predictions and ground truth using Hungarian algorithm.
 
@@ -187,7 +189,7 @@ class HungarianMatcher(nn.Module):
 
 
 def get_cdn_group(
-    batch: Dict[str, Any],
+    batch: dict[str, Any],
     num_classes: int,
     num_queries: int,
     class_embed: torch.Tensor,
@@ -195,7 +197,7 @@ def get_cdn_group(
     cls_noise_ratio: float = 0.5,
     box_noise_scale: float = 1.0,
     training: bool = False,
-) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[Dict[str, Any]]]:
+) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, dict[str, Any] | None]:
     """
     Generate contrastive denoising training group with positive and negative samples from ground truths.
 

ultralytics/models/yolo/classify/train.py

@@ -1,7 +1,9 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from copy import copy
-from typing import Any, Dict, Optional
+from typing import Any
 
 import torch
 
@@ -49,7 +51,7 @@ class ClassificationTrainer(BaseTrainer):
         >>> trainer.train()
     """
 
-    def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[Dict[str, Any]] = None, _callbacks=None):
+    def __init__(self, cfg=DEFAULT_CFG, overrides: dict[str, Any] | None = None, _callbacks=None):
         """
         Initialize a ClassificationTrainer object.
 
@@ -162,7 +164,7 @@ class ClassificationTrainer(BaseTrainer):
                 self.model.transforms = loader.dataset.torch_transforms
         return loader
 
-    def preprocess_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    def preprocess_batch(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         """Preprocess a batch of images and classes."""
         batch["img"] = batch["img"].to(self.device)
         batch["cls"] = batch["cls"].to(self.device)
@@ -185,7 +187,7 @@ class ClassificationTrainer(BaseTrainer):
             self.test_loader, self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
 
-    def label_loss_items(self, loss_items: Optional[torch.Tensor] = None, prefix: str = "train"):
+    def label_loss_items(self, loss_items: torch.Tensor | None = None, prefix: str = "train"):
         """
         Return a loss dict with labelled training loss items tensor.
 
@@ -220,7 +222,7 @@ class ClassificationTrainer(BaseTrainer):
                     self.metrics.pop("fitness", None)
                     self.run_callbacks("on_fit_epoch_end")
 
-    def plot_training_samples(self, batch: Dict[str, torch.Tensor], ni: int):
+    def plot_training_samples(self, batch: dict[str, torch.Tensor], ni: int):
         """
         Plot training samples with their annotations.
 

ultralytics/models/yolo/classify/val.py

@@ -1,7 +1,9 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from pathlib import Path
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any
 
 import torch
 
@@ -85,14 +87,14 @@ class ClassificationValidator(BaseValidator):
         self.targets = []
         self.confusion_matrix = ConfusionMatrix(names=model.names)
 
-    def preprocess(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+    def preprocess(self, batch: dict[str, Any]) -> dict[str, Any]:
         """Preprocess input batch by moving data to device and converting to appropriate dtype."""
         batch["img"] = batch["img"].to(self.device, non_blocking=True)
         batch["img"] = batch["img"].half() if self.args.half else batch["img"].float()
         batch["cls"] = batch["cls"].to(self.device)
         return batch
 
-    def update_metrics(self, preds: torch.Tensor, batch: Dict[str, Any]) -> None:
+    def update_metrics(self, preds: torch.Tensor, batch: dict[str, Any]) -> None:
         """
         Update running metrics with model predictions and batch targets.
 
@@ -131,11 +133,11 @@ class ClassificationValidator(BaseValidator):
         self.metrics.save_dir = self.save_dir
         self.metrics.confusion_matrix = self.confusion_matrix
 
-    def postprocess(self, preds: Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]]) -> torch.Tensor:
+    def postprocess(self, preds: torch.Tensor | list[torch.Tensor] | tuple[torch.Tensor]) -> torch.Tensor:
         """Extract the primary prediction from model output if it's in a list or tuple format."""
         return preds[0] if isinstance(preds, (list, tuple)) else preds
 
-    def get_stats(self) -> Dict[str, float]:
+    def get_stats(self) -> dict[str, float]:
         """Calculate and return a dictionary of metrics by processing targets and predictions."""
         self.metrics.process(self.targets, self.pred)
         return self.metrics.results_dict
@@ -144,7 +146,7 @@ class ClassificationValidator(BaseValidator):
         """Create a ClassificationDataset instance for validation."""
         return ClassificationDataset(root=img_path, args=self.args, augment=False, prefix=self.args.split)
 
-    def get_dataloader(self, dataset_path: Union[Path, str], batch_size: int) -> torch.utils.data.DataLoader:
+    def get_dataloader(self, dataset_path: Path | str, batch_size: int) -> torch.utils.data.DataLoader:
         """
         Build and return a data loader for classification validation.
 
@@ -163,7 +165,7 @@ class ClassificationValidator(BaseValidator):
         pf = "%22s" + "%11.3g" * len(self.metrics.keys)  # print format
         LOGGER.info(pf % ("all", self.metrics.top1, self.metrics.top5))
 
-    def plot_val_samples(self, batch: Dict[str, Any], ni: int) -> None:
+    def plot_val_samples(self, batch: dict[str, Any], ni: int) -> None:
         """
         Plot validation image samples with their ground truth labels.
 
@@ -184,7 +186,7 @@ class ClassificationValidator(BaseValidator):
             on_plot=self.on_plot,
         )
 
-    def plot_predictions(self, batch: Dict[str, Any], preds: torch.Tensor, ni: int) -> None:
+    def plot_predictions(self, batch: dict[str, Any], preds: torch.Tensor, ni: int) -> None:
         """
         Plot images with their predicted class labels and save the visualization.
 

ultralytics/models/yolo/detect/predict.py

@@ -85,7 +85,7 @@ class DetectionPredictor(BasePredictor):
         """Extract object features from the feature maps."""
         import torch
 
-        s = min([x.shape[1] for x in feat_maps])  # find smallest vector length
+        s = min(x.shape[1] for x in feat_maps)  # find shortest vector length
         obj_feats = torch.cat(
             [x.permute(0, 2, 3, 1).reshape(x.shape[0], -1, s, x.shape[1] // s).mean(dim=-1) for x in feat_maps], dim=1
         )  # mean reduce all vectors to same length

ultralytics/models/yolo/detect/train.py

@@ -1,9 +1,11 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import math
 import random
 from copy import copy
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 import numpy as np
 import torch.nn as nn
@@ -51,7 +53,7 @@ class DetectionTrainer(BaseTrainer):
         >>> trainer.train()
     """
 
-    def build_dataset(self, img_path: str, mode: str = "train", batch: Optional[int] = None):
+    def build_dataset(self, img_path: str, mode: str = "train", batch: int | None = None):
         """
         Build YOLO Dataset for training or validation.
 
@@ -89,7 +91,7 @@ class DetectionTrainer(BaseTrainer):
         workers = self.args.workers if mode == "train" else self.args.workers * 2
         return build_dataloader(dataset, batch_size, workers, shuffle, rank)  # return dataloader
 
-    def preprocess_batch(self, batch: Dict) -> Dict:
+    def preprocess_batch(self, batch: dict) -> dict:
         """
         Preprocess a batch of images by scaling and converting to float.
 
@@ -127,7 +129,7 @@ class DetectionTrainer(BaseTrainer):
         self.model.args = self.args  # attach hyperparameters to model
         # TODO: self.model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc
 
-    def get_model(self, cfg: Optional[str] = None, weights: Optional[str] = None, verbose: bool = True):
+    def get_model(self, cfg: str | None = None, weights: str | None = None, verbose: bool = True):
         """
         Return a YOLO detection model.
 
@@ -151,7 +153,7 @@ class DetectionTrainer(BaseTrainer):
             self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
 
-    def label_loss_items(self, loss_items: Optional[List[float]] = None, prefix: str = "train"):
+    def label_loss_items(self, loss_items: list[float] | None = None, prefix: str = "train"):
         """
         Return a loss dict with labeled training loss items tensor.
 
@@ -179,7 +181,7 @@ class DetectionTrainer(BaseTrainer):
             "Size",
         )
 
-    def plot_training_samples(self, batch: Dict[str, Any], ni: int) -> None:
+    def plot_training_samples(self, batch: dict[str, Any], ni: int) -> None:
         """
         Plot training samples with their annotations.
 

ultralytics/models/yolo/detect/val.py

@@ -1,8 +1,10 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import os
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
 import torch
@@ -59,7 +61,7 @@ class DetectionValidator(BaseValidator):
         self.niou = self.iouv.numel()
         self.metrics = DetMetrics()
 
-    def preprocess(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+    def preprocess(self, batch: dict[str, Any]) -> dict[str, Any]:
         """
         Preprocess batch of images for YOLO validation.
 
@@ -104,7 +106,7 @@ class DetectionValidator(BaseValidator):
         """Return a formatted string summarizing class metrics of YOLO model."""
         return ("%22s" + "%11s" * 6) % ("Class", "Images", "Instances", "Box(P", "R", "mAP50", "mAP50-95)")
 
-    def postprocess(self, preds: torch.Tensor) -> List[Dict[str, torch.Tensor]]:
+    def postprocess(self, preds: torch.Tensor) -> list[dict[str, torch.Tensor]]:
         """
         Apply Non-maximum suppression to prediction outputs.
 
@@ -128,7 +130,7 @@ class DetectionValidator(BaseValidator):
         )
         return [{"bboxes": x[:, :4], "conf": x[:, 4], "cls": x[:, 5], "extra": x[:, 6:]} for x in outputs]
 
-    def _prepare_batch(self, si: int, batch: Dict[str, Any]) -> Dict[str, Any]:
+    def _prepare_batch(self, si: int, batch: dict[str, Any]) -> dict[str, Any]:
         """
         Prepare a batch of images and annotations for validation.
 
@@ -156,7 +158,7 @@ class DetectionValidator(BaseValidator):
             "im_file": batch["im_file"][si],
         }
 
-    def _prepare_pred(self, pred: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    def _prepare_pred(self, pred: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         """
         Prepare predictions for evaluation against ground truth.
 
@@ -170,7 +172,7 @@ class DetectionValidator(BaseValidator):
             pred["cls"] *= 0
         return pred
 
-    def update_metrics(self, preds: List[Dict[str, torch.Tensor]], batch: Dict[str, Any]) -> None:
+    def update_metrics(self, preds: list[dict[str, torch.Tensor]], batch: dict[str, Any]) -> None:
         """
         Update metrics with new predictions and ground truth.
 
@@ -225,7 +227,7 @@ class DetectionValidator(BaseValidator):
         self.metrics.confusion_matrix = self.confusion_matrix
         self.metrics.save_dir = self.save_dir
 
-    def get_stats(self) -> Dict[str, Any]:
+    def get_stats(self) -> dict[str, Any]:
         """
         Calculate and return metrics statistics.
 
@@ -256,7 +258,7 @@ class DetectionValidator(BaseValidator):
                     )
                 )
 
-    def _process_batch(self, preds: Dict[str, torch.Tensor], batch: Dict[str, Any]) -> Dict[str, np.ndarray]:
+    def _process_batch(self, preds: dict[str, torch.Tensor], batch: dict[str, Any]) -> dict[str, np.ndarray]:
         """
         Return correct prediction matrix.
 
@@ -272,7 +274,7 @@ class DetectionValidator(BaseValidator):
         iou = box_iou(batch["bboxes"], preds["bboxes"])
         return {"tp": self.match_predictions(preds["cls"], batch["cls"], iou).cpu().numpy()}
 
-    def build_dataset(self, img_path: str, mode: str = "val", batch: Optional[int] = None) -> torch.utils.data.Dataset:
+    def build_dataset(self, img_path: str, mode: str = "val", batch: int | None = None) -> torch.utils.data.Dataset:
         """
         Build YOLO Dataset.
 
@@ -300,7 +302,7 @@ class DetectionValidator(BaseValidator):
         dataset = self.build_dataset(dataset_path, batch=batch_size, mode="val")
         return build_dataloader(dataset, batch_size, self.args.workers, shuffle=False, rank=-1)  # return dataloader
 
-    def plot_val_samples(self, batch: Dict[str, Any], ni: int) -> None:
+    def plot_val_samples(self, batch: dict[str, Any], ni: int) -> None:
         """
         Plot validation image samples.
 
@@ -317,7 +319,7 @@ class DetectionValidator(BaseValidator):
         )
 
     def plot_predictions(
-        self, batch: Dict[str, Any], preds: List[Dict[str, torch.Tensor]], ni: int, max_det: Optional[int] = None
+        self, batch: dict[str, Any], preds: list[dict[str, torch.Tensor]], ni: int, max_det: int | None = None
     ) -> None:
         """
         Plot predicted bounding boxes on input images and save the result.
@@ -345,7 +347,7 @@ class DetectionValidator(BaseValidator):
             on_plot=self.on_plot,
         )  # pred
 
-    def save_one_txt(self, predn: Dict[str, torch.Tensor], save_conf: bool, shape: Tuple[int, int], file: Path) -> None:
+    def save_one_txt(self, predn: dict[str, 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.
 
@@ -364,7 +366,7 @@ class DetectionValidator(BaseValidator):
             boxes=torch.cat([predn["bboxes"], predn["conf"].unsqueeze(-1), predn["cls"].unsqueeze(-1)], dim=1),
         ).save_txt(file, save_conf=save_conf)
 
-    def pred_to_json(self, predn: Dict[str, torch.Tensor], pbatch: Dict[str, Any]) -> None:
+    def pred_to_json(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> None:
         """
         Serialize YOLO predictions to COCO json format.
 
@@ -398,7 +400,7 @@ class DetectionValidator(BaseValidator):
                 }
             )
 
-    def scale_preds(self, predn: Dict[str, torch.Tensor], pbatch: Dict[str, Any]) -> Dict[str, torch.Tensor]:
+    def scale_preds(self, predn: dict[str, torch.Tensor], pbatch: dict[str, Any]) -> dict[str, torch.Tensor]:
         """Scales predictions to the original image size."""
         return {
             **predn,
@@ -410,7 +412,7 @@ class DetectionValidator(BaseValidator):
             ),
         }
 
-    def eval_json(self, stats: Dict[str, Any]) -> Dict[str, Any]:
+    def eval_json(self, stats: dict[str, Any]) -> dict[str, Any]:
         """
         Evaluate YOLO output in JSON format and return performance statistics.
 
@@ -430,12 +432,12 @@ class DetectionValidator(BaseValidator):
 
     def coco_evaluate(
         self,
-        stats: Dict[str, Any],
+        stats: dict[str, Any],
         pred_json: str,
         anno_json: str,
-        iou_types: Union[str, List[str]] = "bbox",
-        suffix: Union[str, List[str]] = "Box",
-    ) -> Dict[str, Any]:
+        iou_types: str | list[str] = "bbox",
+        suffix: str | list[str] = "Box",
+    ) -> dict[str, Any]:
         """
         Evaluate COCO/LVIS metrics using faster-coco-eval library.
 

ultralytics/models/yolo/model.py

@@ -1,7 +1,9 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 import torch
 
@@ -49,7 +51,7 @@ class YOLO(Model):
         >>> model = YOLO("yolo11n.yaml")
     """
 
-    def __init__(self, model: Union[str, Path] = "yolo11n.pt", task: Optional[str] = None, verbose: bool = False):
+    def __init__(self, model: str | Path = "yolo11n.pt", task: str | None = None, verbose: bool = False):
         """
         Initialize a YOLO model.
 
@@ -87,7 +89,7 @@ class YOLO(Model):
                 self.__dict__ = new_instance.__dict__
 
     @property
-    def task_map(self) -> Dict[str, Dict[str, Any]]:
+    def task_map(self) -> dict[str, dict[str, Any]]:
         """Map head to model, trainer, validator, and predictor classes."""
         return {
             "classify": {
@@ -149,7 +151,7 @@ class YOLOWorld(Model):
         >>> model.set_classes(["person", "car", "bicycle"])
     """
 
-    def __init__(self, model: Union[str, Path] = "yolov8s-world.pt", verbose: bool = False) -> None:
+    def __init__(self, model: str | Path = "yolov8s-world.pt", verbose: bool = False) -> None:
         """
         Initialize YOLOv8-World model with a pre-trained model file.
 
@@ -167,7 +169,7 @@ class YOLOWorld(Model):
             self.model.names = YAML.load(ROOT / "cfg/datasets/coco8.yaml").get("names")
 
     @property
-    def task_map(self) -> Dict[str, Dict[str, Any]]:
+    def task_map(self) -> dict[str, dict[str, Any]]:
         """Map head to model, validator, and predictor classes."""
         return {
             "detect": {
@@ -178,7 +180,7 @@ class YOLOWorld(Model):
             }
         }
 
-    def set_classes(self, classes: List[str]) -> None:
+    def set_classes(self, classes: list[str]) -> None:
         """
         Set the model's class names for detection.
 
@@ -232,9 +234,7 @@ class YOLOE(Model):
         >>> 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:
+    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.
 
@@ -246,7 +246,7 @@ class YOLOE(Model):
         super().__init__(model=model, task=task, verbose=verbose)
 
     @property
-    def task_map(self) -> Dict[str, Dict[str, Any]]:
+    def task_map(self) -> dict[str, dict[str, Any]]:
         """Map head to model, validator, and predictor classes."""
         return {
             "detect": {
@@ -291,7 +291,7 @@ class YOLOE(Model):
         assert isinstance(self.model, YOLOEModel)
         return self.model.get_visual_pe(img, visual)
 
-    def set_vocab(self, vocab: List[str], names: List[str]) -> None:
+    def set_vocab(self, vocab: list[str], names: list[str]) -> None:
         """
         Set vocabulary and class names for the YOLOE model.
 
@@ -317,7 +317,7 @@ class YOLOE(Model):
         assert isinstance(self.model, YOLOEModel)
         return self.model.get_vocab(names)
 
-    def set_classes(self, classes: List[str], embeddings: Optional[torch.Tensor] = None) -> None:
+    def set_classes(self, classes: list[str], embeddings: torch.Tensor | None = None) -> None:
         """
         Set the model's class names and embeddings for detection.
 
@@ -341,7 +341,7 @@ class YOLOE(Model):
         self,
         validator=None,
         load_vp: bool = False,
-        refer_data: Optional[str] = None,
+        refer_data: str | None = None,
         **kwargs,
     ):
         """
@@ -368,7 +368,7 @@ class YOLOE(Model):
         self,
         source=None,
         stream: bool = False,
-        visual_prompts: Dict[str, List] = {},
+        visual_prompts: dict[str, list] = {},
         refer_image=None,
         predictor=yolo.yoloe.YOLOEVPDetectPredictor,
         **kwargs,

ultralytics/models/yolo/obb/train.py

@@ -1,8 +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, List, Optional, Union
+from typing import Any
 
 from ultralytics.models import yolo
 from ultralytics.nn.tasks import OBBModel
@@ -31,7 +33,7 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
         >>> trainer.train()
     """
 
-    def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[dict] = None, _callbacks: Optional[List[Any]] = None):
+    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.
 
@@ -57,7 +59,7 @@ class OBBTrainer(yolo.detect.DetectionTrainer):
         super().__init__(cfg, overrides, _callbacks)
 
     def get_model(
-        self, cfg: Optional[Union[str, dict]] = None, weights: Optional[Union[str, Path]] = None, verbose: bool = True
+        self, cfg: str | dict | None = None, weights: str | Path | None = None, verbose: bool = True
     ) -> OBBModel:
         """
         Return OBBModel initialized with specified config and weights.