单目3D初始代码

ultralytics/utils/autobatch.py

@@ -0,0 +1,119 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""Functions for estimating the best YOLO batch size to use a fraction of the available CUDA memory in PyTorch."""
+
+from __future__ import annotations
+
+import os
+from copy import deepcopy
+
+import numpy as np
+import torch
+
+from ultralytics.utils import DEFAULT_CFG, LOGGER, colorstr
+from ultralytics.utils.torch_utils import autocast, profile_ops
+
+
+def check_train_batch_size(
+    model: torch.nn.Module,
+    imgsz: int = 640,
+    amp: bool = True,
+    batch: int | float = -1,
+    max_num_obj: int = 1,
+) -> int:
+    """Compute optimal YOLO training batch size using the autobatch() function.
+
+    Args:
+        model (torch.nn.Module): YOLO model to check batch size for.
+        imgsz (int, optional): Image size used for training.
+        amp (bool, optional): Use automatic mixed precision if True.
+        batch (int | float, optional): Fraction of GPU memory to use. If -1, use default.
+        max_num_obj (int, optional): The maximum number of objects from dataset.
+
+    Returns:
+        (int): Optimal batch size computed using the autobatch() function.
+
+    Notes:
+        If 0.0 < batch < 1.0, it's used as the fraction of GPU memory to use.
+        Otherwise, a default fraction of 0.6 is used.
+    """
+    with autocast(enabled=amp):
+        return autobatch(
+            deepcopy(model).train(), imgsz, fraction=batch if 0.0 < batch < 1.0 else 0.6, max_num_obj=max_num_obj
+        )
+
+
+def autobatch(
+    model: torch.nn.Module,
+    imgsz: int = 640,
+    fraction: float = 0.60,
+    batch_size: int = DEFAULT_CFG.batch,
+    max_num_obj: int = 1,
+) -> int:
+    """Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.
+
+    Args:
+        model (torch.nn.Module): YOLO model to compute batch size for.
+        imgsz (int, optional): The image size used as input for the YOLO model.
+        fraction (float, optional): The fraction of available CUDA memory to use.
+        batch_size (int, optional): The default batch size to use if an error is detected.
+        max_num_obj (int, optional): The maximum number of objects from dataset.
+
+    Returns:
+        (int): The optimal batch size.
+    """
+    # Check device
+    prefix = colorstr("AutoBatch: ")
+    LOGGER.info(f"{prefix}Computing optimal batch size for imgsz={imgsz} at {fraction * 100}% CUDA memory utilization.")
+    device = next(model.parameters()).device  # get model device
+    if device.type in {"cpu", "mps"}:
+        LOGGER.warning(f"{prefix}intended for CUDA devices, using default batch-size {batch_size}")
+        return batch_size
+    if torch.backends.cudnn.benchmark:
+        LOGGER.warning(f"{prefix}Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}")
+        return batch_size
+
+    # Inspect CUDA memory
+    gb = 1 << 30  # bytes to GiB (1024 ** 3)
+    d = f"CUDA:{os.getenv('CUDA_VISIBLE_DEVICES', '0').strip()[0]}"  # 'CUDA:0'
+    properties = torch.cuda.get_device_properties(device)  # device properties
+    t = properties.total_memory / gb  # GiB total
+    r = torch.cuda.memory_reserved(device) / gb  # GiB reserved
+    a = torch.cuda.memory_allocated(device) / gb  # GiB allocated
+    f = t - (r + a)  # GiB free
+    LOGGER.info(f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free")
+
+    # Profile batch sizes
+    batch_sizes = [1, 2, 4, 8, 16] if t < 16 else [1, 2, 4, 8, 16, 32, 64]
+    ch = model.yaml.get("channels", 3)
+    try:
+        img = [torch.empty(b, ch, imgsz, imgsz) for b in batch_sizes]
+        results = profile_ops(img, model, n=1, device=device, max_num_obj=max_num_obj)
+
+        # Fit a solution
+        xy = [
+            [x, y[2]]
+            for i, (x, y) in enumerate(zip(batch_sizes, results))
+            if y  # valid result
+            and isinstance(y[2], (int, float))  # is numeric
+            and 0 < y[2] < t  # between 0 and GPU limit
+            and (i == 0 or not results[i - 1] or y[2] > results[i - 1][2])  # first item or increasing memory
+        ]
+        fit_x, fit_y = zip(*xy) if xy else ([], [])
+        p = np.polyfit(fit_x, fit_y, deg=1)  # first-degree polynomial fit in log space
+        b = int((round(f * fraction) - p[1]) / p[0])  # y intercept (optimal batch size)
+        if None in results:  # some sizes failed
+            i = results.index(None)  # first fail index
+            if b >= batch_sizes[i]:  # y intercept above failure point
+                b = batch_sizes[max(i - 1, 0)]  # select prior safe point
+        if b < 1 or b > 1024:  # b outside of safe range
+            LOGGER.warning(f"{prefix}batch={b} outside safe range, using default batch-size {batch_size}.")
+            b = batch_size
+
+        fraction = (np.polyval(p, b) + r + a) / t  # predicted fraction
+        LOGGER.info(f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅")
+        return b
+    except Exception as e:
+        LOGGER.warning(f"{prefix}error detected: {e},  using default batch-size {batch_size}.")
+        return batch_size
+    finally:
+        torch.cuda.empty_cache()

ultralytics/utils/autodevice.py

@@ -0,0 +1,208 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import random
+from typing import Any
+
+from ultralytics.utils import LOGGER
+from ultralytics.utils.checks import check_requirements
+
+
+class GPUInfo:
+    """Manages NVIDIA GPU information via pynvml with robust error handling.
+
+    Provides methods to query detailed GPU statistics (utilization, memory, temp, power) and select the most idle GPUs
+    based on configurable criteria. It safely handles the absence or initialization failure of the pynvml library by
+    logging warnings and disabling related features, preventing application crashes.
+
+    Includes fallback logic using `torch.cuda` for basic device counting if NVML is unavailable during GPU
+    selection. Manages NVML initialization and shutdown internally.
+
+    Attributes:
+        pynvml (module | None): The `pynvml` module if successfully imported and initialized, otherwise `None`.
+        nvml_available (bool): Indicates if `pynvml` is ready for use. True if import and `nvmlInit()` succeeded, False
+            otherwise.
+        gpu_stats (list[dict[str, Any]]): A list of dictionaries, each holding stats for one GPU, populated on
+        initialization and by `refresh_stats()`. Keys include: 'index', 'name', 'utilization' (%), 'memory_used' (MiB),
+            'memory_total' (MiB), 'memory_free' (MiB), 'temperature' (C), 'power_draw' (W), 'power_limit' (W or 'N/A').
+            Empty if NVML is unavailable or queries fail.
+
+    Methods:
+        refresh_stats: Refresh the internal gpu_stats list by querying NVML.
+        print_status: Print GPU status in a compact table format using current stats.
+        select_idle_gpu: Select the most idle GPUs based on utilization and free memory.
+        shutdown: Shut down NVML if it was initialized.
+
+    Examples:
+        Initialize GPUInfo and print status
+        >>> gpu_info = GPUInfo()
+        >>> gpu_info.print_status()
+
+        Select idle GPUs with minimum memory requirements
+        >>> selected = gpu_info.select_idle_gpu(count=2, min_memory_fraction=0.2)
+        >>> print(f"Selected GPU indices: {selected}")
+    """
+
+    def __init__(self):
+        """Initialize GPUInfo, attempting to import and initialize pynvml."""
+        self.pynvml: Any | None = None
+        self.nvml_available: bool = False
+        self.gpu_stats: list[dict[str, Any]] = []
+
+        try:
+            check_requirements("nvidia-ml-py>=12.0.0")
+            self.pynvml = __import__("pynvml")
+            self.pynvml.nvmlInit()
+            self.nvml_available = True
+            self.refresh_stats()
+        except Exception as e:
+            LOGGER.warning(f"Failed to initialize pynvml, GPU stats disabled: {e}")
+
+    def __del__(self):
+        """Ensure NVML is shut down when the object is garbage collected."""
+        self.shutdown()
+
+    def shutdown(self):
+        """Shut down NVML if it was initialized."""
+        if self.nvml_available and self.pynvml:
+            try:
+                self.pynvml.nvmlShutdown()
+            except Exception:
+                pass
+            self.nvml_available = False
+
+    def refresh_stats(self):
+        """Refresh the internal gpu_stats list by querying NVML."""
+        self.gpu_stats = []
+        if not self.nvml_available or not self.pynvml:
+            return
+
+        try:
+            device_count = self.pynvml.nvmlDeviceGetCount()
+            self.gpu_stats.extend(self._get_device_stats(i) for i in range(device_count))
+        except Exception as e:
+            LOGGER.warning(f"Error during device query: {e}")
+            self.gpu_stats = []
+
+    def _get_device_stats(self, index: int) -> dict[str, Any]:
+        """Get stats for a single GPU device."""
+        handle = self.pynvml.nvmlDeviceGetHandleByIndex(index)
+        memory = self.pynvml.nvmlDeviceGetMemoryInfo(handle)
+        util = self.pynvml.nvmlDeviceGetUtilizationRates(handle)
+
+        def safe_get(func, *args, default=-1, divisor=1):
+            try:
+                val = func(*args)
+                return val // divisor if divisor != 1 and isinstance(val, (int, float)) else val
+            except Exception:
+                return default
+
+        temp_type = getattr(self.pynvml, "NVML_TEMPERATURE_GPU", -1)
+
+        return {
+            "index": index,
+            "name": self.pynvml.nvmlDeviceGetName(handle),
+            "utilization": util.gpu if util else -1,
+            "memory_used": memory.used >> 20 if memory else -1,  # Convert bytes to MiB
+            "memory_total": memory.total >> 20 if memory else -1,
+            "memory_free": memory.free >> 20 if memory else -1,
+            "temperature": safe_get(self.pynvml.nvmlDeviceGetTemperature, handle, temp_type),
+            "power_draw": safe_get(self.pynvml.nvmlDeviceGetPowerUsage, handle, divisor=1000),  # Convert mW to W
+            "power_limit": safe_get(self.pynvml.nvmlDeviceGetEnforcedPowerLimit, handle, divisor=1000),
+        }
+
+    def print_status(self):
+        """Print GPU status in a compact table format using current stats."""
+        self.refresh_stats()
+        if not self.gpu_stats:
+            LOGGER.warning("No GPU stats available.")
+            return
+
+        stats = self.gpu_stats
+        name_len = max(len(gpu.get("name", "N/A")) for gpu in stats)
+        hdr = f"{'Idx':<3} {'Name':<{name_len}} {'Util':>6} {'Mem (MiB)':>15} {'Temp':>5} {'Pwr (W)':>10}"
+        LOGGER.info(f"\n--- GPU Status ---\n{hdr}\n{'-' * len(hdr)}")
+
+        for gpu in stats:
+            u = f"{gpu['utilization']:>5}%" if gpu["utilization"] >= 0 else " N/A "
+            m = f"{gpu['memory_used']:>6}/{gpu['memory_total']:<6}" if gpu["memory_used"] >= 0 else " N/A / N/A "
+            t = f"{gpu['temperature']}C" if gpu["temperature"] >= 0 else " N/A "
+            p = f"{gpu['power_draw']:>3}/{gpu['power_limit']:<3}" if gpu["power_draw"] >= 0 else " N/A "
+
+            LOGGER.info(f"{gpu.get('index'):<3d} {gpu.get('name', 'N/A'):<{name_len}} {u:>6} {m:>15} {t:>5} {p:>10}")
+
+        LOGGER.info(f"{'-' * len(hdr)}\n")
+
+    def select_idle_gpu(
+        self, count: int = 1, min_memory_fraction: float = 0, min_util_fraction: float = 0
+    ) -> list[int]:
+        """Select the most idle GPUs based on utilization and free memory.
+
+        Args:
+            count (int): The number of idle GPUs to select.
+            min_memory_fraction (float): Minimum free memory required as a fraction of total memory.
+            min_util_fraction (float): Minimum free utilization rate required from 0.0 - 1.0.
+
+        Returns:
+            (list[int]): Indices of the selected GPUs, sorted by idleness (lowest utilization first).
+
+        Notes:
+             Returns fewer than 'count' if not enough qualify or exist.
+             Returns empty list if NVML stats are unavailable or no GPUs meet the criteria.
+        """
+        assert min_memory_fraction <= 1.0, f"min_memory_fraction must be <= 1.0, got {min_memory_fraction}"
+        assert min_util_fraction <= 1.0, f"min_util_fraction must be <= 1.0, got {min_util_fraction}"
+        criteria = (
+            f"free memory >= {min_memory_fraction * 100:.1f}% and free utilization >= {min_util_fraction * 100:.1f}%"
+        )
+        LOGGER.info(f"Searching for {count} idle GPUs with {criteria}...")
+
+        if count <= 0:
+            return []
+
+        self.refresh_stats()
+        if not self.gpu_stats:
+            LOGGER.warning("NVML stats unavailable.")
+            return []
+
+        # Filter and sort eligible GPUs
+        eligible_gpus = [
+            gpu
+            for gpu in self.gpu_stats
+            if gpu.get("memory_free", 0) / gpu.get("memory_total", 1) >= min_memory_fraction
+            and (100 - gpu.get("utilization", 100)) >= min_util_fraction * 100
+        ]
+        # Random tiebreaker prevents race conditions when multiple processes start simultaneously
+        # and all GPUs appear equally idle (same utilization and free memory)
+        eligible_gpus.sort(key=lambda x: (x.get("utilization", 101), -x.get("memory_free", 0), random.random()))
+
+        # Select top 'count' indices
+        selected = [gpu["index"] for gpu in eligible_gpus[:count]]
+
+        if selected:
+            if len(selected) < count:
+                LOGGER.warning(f"Requested {count} GPUs but only {len(selected)} met the idle criteria.")
+            LOGGER.info(f"Selected idle CUDA devices {selected}")
+        else:
+            LOGGER.warning(f"No GPUs met criteria ({criteria}).")
+
+        return selected
+
+
+if __name__ == "__main__":
+    required_free_mem_fraction = 0.2  # Require 20% free VRAM
+    required_free_util_fraction = 0.2  # Require 20% free utilization
+    num_gpus_to_select = 1
+
+    gpu_info = GPUInfo()
+    gpu_info.print_status()
+
+    if selected := gpu_info.select_idle_gpu(
+        count=num_gpus_to_select,
+        min_memory_fraction=required_free_mem_fraction,
+        min_util_fraction=required_free_util_fraction,
+    ):
+        print(f"\n==> Using selected GPU indices: {selected}")
+        devices = [f"cuda:{idx}" for idx in selected]
+        print(f"    Target devices: {devices}")

ultralytics/utils/benchmarks.py

@@ -0,0 +1,718 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""
+Benchmark YOLO model formats for speed and accuracy.
+
+Usage:
+    from ultralytics.utils.benchmarks import ProfileModels, benchmark
+    ProfileModels(['yolo26n.yaml', 'yolov8s.yaml']).run()
+    benchmark(model='yolo26n.pt', imgsz=160)
+
+Format                  | `format=argument`         | Model
+---                     | ---                       | ---
+PyTorch                 | -                         | yolo26n.pt
+TorchScript             | `torchscript`             | yolo26n.torchscript
+ONNX                    | `onnx`                    | yolo26n.onnx
+OpenVINO                | `openvino`                | yolo26n_openvino_model/
+TensorRT                | `engine`                  | yolo26n.engine
+CoreML                  | `coreml`                  | yolo26n.mlpackage
+TensorFlow SavedModel   | `saved_model`             | yolo26n_saved_model/
+TensorFlow GraphDef     | `pb`                      | yolo26n.pb
+TensorFlow Lite         | `tflite`                  | yolo26n.tflite
+TensorFlow Edge TPU     | `edgetpu`                 | yolo26n_edgetpu.tflite
+TensorFlow.js           | `tfjs`                    | yolo26n_web_model/
+PaddlePaddle            | `paddle`                  | yolo26n_paddle_model/
+MNN                     | `mnn`                     | yolo26n.mnn
+NCNN                    | `ncnn`                    | yolo26n_ncnn_model/
+IMX                     | `imx`                     | yolo26n_imx_model/
+RKNN                    | `rknn`                    | yolo26n_rknn_model/
+ExecuTorch              | `executorch`              | yolo26n_executorch_model/
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+import platform
+import re
+import shutil
+import time
+from copy import deepcopy
+from pathlib import Path
+
+import numpy as np
+import torch.cuda
+
+from ultralytics import YOLO, YOLOWorld
+from ultralytics.cfg import TASK2DATA, TASK2METRIC
+from ultralytics.engine.exporter import export_formats
+from ultralytics.utils import ARM64, ASSETS, ASSETS_URL, IS_JETSON, LINUX, LOGGER, MACOS, TQDM, WEIGHTS_DIR, YAML
+from ultralytics.utils.checks import IS_PYTHON_3_13, check_imgsz, check_requirements, check_yolo, is_rockchip
+from ultralytics.utils.downloads import safe_download
+from ultralytics.utils.files import file_size
+from ultralytics.utils.torch_utils import get_cpu_info, select_device
+
+
+def benchmark(
+    model=WEIGHTS_DIR / "yolo26n.pt",
+    data=None,
+    imgsz=160,
+    half=False,
+    int8=False,
+    device="cpu",
+    verbose=False,
+    eps=1e-3,
+    format="",
+    **kwargs,
+):
+    """Benchmark a YOLO model across different formats for speed and accuracy.
+
+    Args:
+        model (str | Path): Path to the model file or directory.
+        data (str | None): Dataset to evaluate on, inherited from TASK2DATA if not passed.
+        imgsz (int): Image size for the benchmark.
+        half (bool): Use half-precision for the model if True.
+        int8 (bool): Use int8-precision for the model if True.
+        device (str): Device to run the benchmark on, either 'cpu' or 'cuda'.
+        verbose (bool | float): If True or a float, assert benchmarks pass with given metric.
+        eps (float): Epsilon value for divide by zero prevention.
+        format (str): Export format for benchmarking. If not supplied all formats are benchmarked.
+        **kwargs (Any): Additional keyword arguments for exporter.
+
+    Returns:
+        (polars.DataFrame): A Polars DataFrame with benchmark results for each format, including file size, metric, and
+            inference time.
+
+    Examples:
+        Benchmark a YOLO model with default settings:
+        >>> from ultralytics.utils.benchmarks import benchmark
+        >>> benchmark(model="yolo26n.pt", imgsz=640)
+    """
+    imgsz = check_imgsz(imgsz)
+    assert imgsz[0] == imgsz[1] if isinstance(imgsz, list) else True, "benchmark() only supports square imgsz."
+
+    import polars as pl  # scope for faster 'import ultralytics'
+
+    pl.Config.set_tbl_cols(-1)  # Show all columns
+    pl.Config.set_tbl_rows(-1)  # Show all rows
+    pl.Config.set_tbl_width_chars(-1)  # No width limit
+    pl.Config.set_tbl_hide_column_data_types(True)  # Hide data types
+    pl.Config.set_tbl_hide_dataframe_shape(True)  # Hide shape info
+    pl.Config.set_tbl_formatting("ASCII_BORDERS_ONLY_CONDENSED")
+
+    device = select_device(device, verbose=False)
+    if isinstance(model, (str, Path)):
+        model = YOLO(model)
+    data = data or TASK2DATA[model.task]  # task to dataset, i.e. coco8.yaml for task=detect
+    key = TASK2METRIC[model.task]  # task to metric, i.e. metrics/mAP50-95(B) for task=detect
+
+    y = []
+    t0 = time.time()
+
+    format_arg = format.lower()
+    if format_arg:
+        formats = frozenset(export_formats()["Argument"])
+        assert format in formats, f"Expected format to be one of {formats}, but got '{format_arg}'."
+    for name, format, suffix, cpu, gpu, _ in zip(*export_formats().values()):
+        emoji, filename = "❌", None  # export defaults
+        try:
+            if format_arg and format_arg != format:
+                continue
+
+            # Checks
+            if format == "pb":
+                assert model.task != "obb", "TensorFlow GraphDef not supported for OBB task"
+            elif format == "edgetpu":
+                assert LINUX and not ARM64, "Edge TPU export only supported on non-aarch64 Linux"
+            elif format in {"coreml", "tfjs"}:
+                assert MACOS or (LINUX and not ARM64), (
+                    "CoreML and TF.js export only supported on macOS and non-aarch64 Linux"
+                )
+            if format == "coreml":
+                assert not IS_PYTHON_3_13, "CoreML not supported on Python 3.13"
+            if format in {"saved_model", "pb", "tflite", "edgetpu", "tfjs"}:
+                assert not isinstance(model, YOLOWorld), "YOLOWorldv2 TensorFlow exports not supported by onnx2tf yet"
+                # assert not IS_PYTHON_MINIMUM_3_12, "TFLite exports not supported on Python>=3.12 yet"
+            if format == "paddle":
+                assert not isinstance(model, YOLOWorld), "YOLOWorldv2 Paddle exports not supported yet"
+                assert model.task != "obb", "Paddle OBB bug https://github.com/PaddlePaddle/Paddle/issues/72024"
+                assert (LINUX and not IS_JETSON) or MACOS, "Windows and Jetson Paddle exports not supported yet"
+            if format == "mnn":
+                assert not isinstance(model, YOLOWorld), "YOLOWorldv2 MNN exports not supported yet"
+            if format == "ncnn":
+                assert not isinstance(model, YOLOWorld), "YOLOWorldv2 NCNN exports not supported yet"
+            if format == "imx":
+                assert not isinstance(model, YOLOWorld), "YOLOWorldv2 IMX exports not supported"
+                assert model.task in {"detect", "classify", "pose", "segment"}, (
+                    "IMX export is only supported for detection, classification, pose estimation and segmentation tasks"
+                )
+                assert "C2f" in model.__str__(), "IMX only supported for YOLOv8n and YOLO11n"
+            if format == "rknn":
+                assert not isinstance(model, YOLOWorld), "YOLOWorldv2 RKNN exports not supported yet"
+                assert LINUX, "RKNN only supported on Linux"
+                assert not is_rockchip(), "RKNN Inference only supported on Rockchip devices"
+            if format == "executorch":
+                assert not isinstance(model, YOLOWorld), "YOLOWorldv2 ExecuTorch exports not supported yet"
+            if "cpu" in device.type:
+                assert cpu, "inference not supported on CPU"
+            if "cuda" in device.type:
+                assert gpu, "inference not supported on GPU"
+
+            # Export
+            if format == "-":
+                filename = model.pt_path or model.ckpt_path or model.model_name
+                exported_model = deepcopy(model)  # PyTorch format
+            else:
+                filename = deepcopy(model).export(
+                    imgsz=imgsz, format=format, half=half, int8=int8, data=data, device=device, verbose=False, **kwargs
+                )
+                exported_model = YOLO(filename, task=model.task)
+                assert suffix in str(filename), "export failed"
+            emoji = "❎"  # indicates export succeeded
+
+            # Predict
+            assert model.task != "pose" or format != "pb", "GraphDef Pose inference is not supported"
+            assert format not in {"edgetpu", "tfjs"}, "inference not supported"
+            assert format != "coreml" or platform.system() == "Darwin", "inference only supported on macOS>=10.13"
+            exported_model.predict(ASSETS / "bus.jpg", imgsz=imgsz, device=device, half=half, verbose=False)
+
+            # Validate
+            results = exported_model.val(
+                data=data,
+                batch=1,
+                imgsz=imgsz,
+                plots=False,
+                device=device,
+                half=half,
+                int8=int8,
+                verbose=False,
+                conf=0.001,  # all the pre-set benchmark mAP values are based on conf=0.001
+            )
+            metric, speed = results.results_dict[key], results.speed["inference"]
+            fps = round(1000 / (speed + eps), 2)  # frames per second
+            y.append([name, "✅", round(file_size(filename), 1), round(metric, 4), round(speed, 2), fps])
+        except Exception as e:
+            if verbose:
+                assert type(e) is AssertionError, f"Benchmark failure for {name}: {e}"
+            LOGGER.error(f"Benchmark failure for {name}: {e}")
+            y.append([name, emoji, round(file_size(filename), 1), None, None, None])  # mAP, t_inference
+
+    # Print results
+    check_yolo(device=device)  # print system info
+    df = pl.DataFrame(y, schema=["Format", "Status❔", "Size (MB)", key, "Inference time (ms/im)", "FPS"], orient="row")
+    df = df.with_row_index(" ", offset=1)  # add index info
+    df_display = df.with_columns(pl.all().cast(pl.String).fill_null("-"))
+
+    name = model.model_name
+    dt = time.time() - t0
+    legend = "Benchmarks legend:  - ✅ Success  - ❎ Export passed but validation failed  - ❌️ Export failed"
+    s = f"\nBenchmarks complete for {name} on {data} at imgsz={imgsz} ({dt:.2f}s)\n{legend}\n{df_display}\n"
+    LOGGER.info(s)
+    with open("benchmarks.log", "a", errors="ignore", encoding="utf-8") as f:
+        f.write(s)
+
+    if verbose and isinstance(verbose, float):
+        metrics = df[key].to_numpy()  # values to compare to floor
+        floor = verbose  # minimum metric floor to pass, i.e. = 0.29 mAP for YOLOv5n
+        assert all(x > floor for x in metrics if not np.isnan(x)), f"Benchmark failure: metric(s) < floor {floor}"
+
+    return df_display
+
+
+class RF100Benchmark:
+    """Benchmark YOLO model performance on the RF100 dataset collection.
+
+    This class provides functionality to download, process, and evaluate YOLO models on the RF100 datasets.
+
+    Attributes:
+        ds_names (list[str]): Names of datasets used for benchmarking.
+        ds_cfg_list (list[Path]): List of paths to dataset configuration files.
+        rf (Roboflow | None): Roboflow instance for accessing datasets.
+        val_metrics (list[str]): Metrics used for validation.
+
+    Methods:
+        set_key: Set Roboflow API key for accessing datasets.
+        parse_dataset: Parse dataset links and download datasets.
+        fix_yaml: Fix train and validation paths in YAML files.
+        evaluate: Evaluate model performance on validation results.
+    """
+
+    def __init__(self):
+        """Initialize the RF100Benchmark class for benchmarking YOLO model performance on RF100 datasets."""
+        self.ds_names = []
+        self.ds_cfg_list = []
+        self.rf = None
+        self.val_metrics = ["class", "images", "targets", "precision", "recall", "map50", "map95"]
+
+    def set_key(self, api_key: str):
+        """Set Roboflow API key for processing.
+
+        Args:
+            api_key (str): The API key.
+
+        Examples:
+            Set the Roboflow API key for accessing datasets:
+            >>> benchmark = RF100Benchmark()
+            >>> benchmark.set_key("your_roboflow_api_key")
+        """
+        check_requirements("roboflow")
+        from roboflow import Roboflow
+
+        self.rf = Roboflow(api_key=api_key)
+
+    def parse_dataset(self, ds_link_txt: str = "datasets_links.txt"):
+        """Parse dataset links and download datasets.
+
+        Args:
+            ds_link_txt (str): Path to the file containing dataset links.
+
+        Returns:
+            (tuple[list[str], list[Path]]): List of dataset names and list of paths to dataset configuration files.
+
+        Examples:
+            >>> benchmark = RF100Benchmark()
+            >>> benchmark.set_key("api_key")
+            >>> benchmark.parse_dataset("datasets_links.txt")
+        """
+        (shutil.rmtree("rf-100"), os.mkdir("rf-100")) if os.path.exists("rf-100") else os.mkdir("rf-100")
+        os.chdir("rf-100")
+        os.mkdir("ultralytics-benchmarks")
+        safe_download(f"{ASSETS_URL}/datasets_links.txt")
+
+        with open(ds_link_txt, encoding="utf-8") as file:
+            for line in file:
+                try:
+                    _, _url, workspace, project, version = re.split("/+", line.strip())
+                    self.ds_names.append(project)
+                    proj_version = f"{project}-{version}"
+                    if not Path(proj_version).exists():
+                        self.rf.workspace(workspace).project(project).version(version).download("yolov8")
+                    else:
+                        LOGGER.info("Dataset already downloaded.")
+                    self.ds_cfg_list.append(Path.cwd() / proj_version / "data.yaml")
+                except Exception:
+                    continue
+
+        return self.ds_names, self.ds_cfg_list
+
+    @staticmethod
+    def fix_yaml(path: Path):
+        """Fix the train and validation paths in a given YAML file."""
+        yaml_data = YAML.load(path)
+        yaml_data["train"] = "train/images"
+        yaml_data["val"] = "valid/images"
+        YAML.dump(yaml_data, path)
+
+    def evaluate(self, yaml_path: str, val_log_file: str, eval_log_file: str, list_ind: int):
+        """Evaluate model performance on validation results.
+
+        Args:
+            yaml_path (str): Path to the YAML configuration file.
+            val_log_file (str): Path to the validation log file.
+            eval_log_file (str): Path to the evaluation log file.
+            list_ind (int): Index of the current dataset in the list.
+
+        Returns:
+            (float): The mean average precision (mAP) value for the evaluated model.
+
+        Examples:
+            Evaluate a model on a specific dataset
+            >>> benchmark = RF100Benchmark()
+            >>> benchmark.evaluate("path/to/data.yaml", "path/to/val_log.txt", "path/to/eval_log.txt", 0)
+        """
+        skip_symbols = ["🚀", "⚠️", "💡", "❌"]
+        class_names = YAML.load(yaml_path)["names"]
+        with open(val_log_file, encoding="utf-8") as f:
+            lines = f.readlines()
+            eval_lines = []
+            for line in lines:
+                if any(symbol in line for symbol in skip_symbols):
+                    continue
+                entries = line.split(" ")
+                entries = list(filter(lambda val: val != "", entries))
+                entries = [e.strip("\n") for e in entries]
+                eval_lines.extend(
+                    {
+                        "class": entries[0],
+                        "images": entries[1],
+                        "targets": entries[2],
+                        "precision": entries[3],
+                        "recall": entries[4],
+                        "map50": entries[5],
+                        "map95": entries[6],
+                    }
+                    for e in entries
+                    if e in class_names or (e == "all" and "(AP)" not in entries and "(AR)" not in entries)
+                )
+        map_val = 0.0
+        if len(eval_lines) > 1:
+            LOGGER.info("Multiple dicts found")
+            for lst in eval_lines:
+                if lst["class"] == "all":
+                    map_val = lst["map50"]
+        else:
+            LOGGER.info("Single dict found")
+            map_val = next(res["map50"] for res in eval_lines)
+
+        with open(eval_log_file, "a", encoding="utf-8") as f:
+            f.write(f"{self.ds_names[list_ind]}: {map_val}\n")
+
+        return float(map_val)
+
+
+class ProfileModels:
+    """ProfileModels class for profiling different models on ONNX and TensorRT.
+
+    This class profiles the performance of different models, returning results such as model speed and FLOPs.
+
+    Attributes:
+        paths (list[str]): Paths of the models to profile.
+        num_timed_runs (int): Number of timed runs for the profiling.
+        num_warmup_runs (int): Number of warmup runs before profiling.
+        min_time (float): Minimum number of seconds to profile for.
+        imgsz (int): Image size used in the models.
+        half (bool): Flag to indicate whether to use FP16 half-precision for TensorRT profiling.
+        trt (bool): Flag to indicate whether to profile using TensorRT.
+        device (torch.device): Device used for profiling.
+
+    Methods:
+        run: Profile YOLO models for speed and accuracy across various formats.
+        get_files: Get all relevant model files.
+        get_onnx_model_info: Extract metadata from an ONNX model.
+        iterative_sigma_clipping: Apply sigma clipping to remove outliers.
+        profile_tensorrt_model: Profile a TensorRT model.
+        profile_onnx_model: Profile an ONNX model.
+        generate_table_row: Generate a table row with model metrics.
+        generate_results_dict: Generate a dictionary of profiling results.
+        print_table: Print a formatted table of results.
+
+    Examples:
+        Profile models and print results
+        >>> from ultralytics.utils.benchmarks import ProfileModels
+        >>> profiler = ProfileModels(["yolo26n.yaml", "yolov8s.yaml"], imgsz=640)
+        >>> profiler.run()
+    """
+
+    def __init__(
+        self,
+        paths: list[str],
+        num_timed_runs: int = 100,
+        num_warmup_runs: int = 10,
+        min_time: float = 60,
+        imgsz: int = 640,
+        half: bool = True,
+        trt: bool = True,
+        device: torch.device | str | None = None,
+    ):
+        """Initialize the ProfileModels class for profiling models.
+
+        Args:
+            paths (list[str]): List of paths of the models to be profiled.
+            num_timed_runs (int): Number of timed runs for the profiling.
+            num_warmup_runs (int): Number of warmup runs before the actual profiling starts.
+            min_time (float): Minimum time in seconds for profiling a model.
+            imgsz (int): Size of the image used during profiling.
+            half (bool): Flag to indicate whether to use FP16 half-precision for TensorRT profiling.
+            trt (bool): Flag to indicate whether to profile using TensorRT.
+            device (torch.device | str | None): Device used for profiling. If None, it is determined automatically.
+
+        Notes:
+            FP16 'half' argument option removed for ONNX as slower on CPU than FP32.
+        """
+        self.paths = paths
+        self.num_timed_runs = num_timed_runs
+        self.num_warmup_runs = num_warmup_runs
+        self.min_time = min_time
+        self.imgsz = imgsz
+        self.half = half
+        self.trt = trt  # run TensorRT profiling
+        self.device = device if isinstance(device, torch.device) else select_device(device)
+
+    def run(self):
+        """Profile YOLO models for speed and accuracy across various formats including ONNX and TensorRT.
+
+        Returns:
+            (list[dict]): List of dictionaries containing profiling results for each model.
+
+        Examples:
+            Profile models and print results
+            >>> from ultralytics.utils.benchmarks import ProfileModels
+            >>> profiler = ProfileModels(["yolo26n.yaml", "yolo11s.yaml"])
+            >>> results = profiler.run()
+        """
+        files = self.get_files()
+
+        if not files:
+            LOGGER.warning("No matching *.pt or *.onnx files found.")
+            return []
+
+        table_rows = []
+        output = []
+        for file in files:
+            engine_file = file.with_suffix(".engine")
+            if file.suffix in {".pt", ".yaml", ".yml"}:
+                model = YOLO(str(file))
+                model.fuse()  # to report correct params and GFLOPs in model.info()
+                model_info = model.info(imgsz=self.imgsz)
+                if self.trt and self.device.type != "cpu" and not engine_file.is_file():
+                    engine_file = model.export(
+                        format="engine",
+                        half=self.half,
+                        imgsz=self.imgsz,
+                        device=self.device,
+                        verbose=False,
+                    )
+                onnx_file = model.export(
+                    format="onnx",
+                    imgsz=self.imgsz,
+                    device=self.device,
+                    verbose=False,
+                )
+            elif file.suffix == ".onnx":
+                model_info = self.get_onnx_model_info(file)
+                onnx_file = file
+            else:
+                continue
+
+            t_engine = self.profile_tensorrt_model(str(engine_file))
+            t_onnx = self.profile_onnx_model(str(onnx_file))
+            table_rows.append(self.generate_table_row(file.stem, t_onnx, t_engine, model_info))
+            output.append(self.generate_results_dict(file.stem, t_onnx, t_engine, model_info))
+
+        self.print_table(table_rows)
+        return output
+
+    def get_files(self):
+        """Return a list of paths for all relevant model files given by the user.
+
+        Returns:
+            (list[Path]): List of Path objects for the model files.
+        """
+        files = []
+        for path in self.paths:
+            path = Path(path)
+            if path.is_dir():
+                extensions = ["*.pt", "*.onnx", "*.yaml"]
+                files.extend([file for ext in extensions for file in glob.glob(str(path / ext))])
+            elif path.suffix in {".pt", ".yaml", ".yml"}:  # add non-existing
+                files.append(str(path))
+            else:
+                files.extend(glob.glob(str(path)))
+
+        LOGGER.info(f"Profiling: {sorted(files)}")
+        return [Path(file) for file in sorted(files)]
+
+    @staticmethod
+    def get_onnx_model_info(onnx_file: str):
+        """Extract metadata from an ONNX model file including layers, parameters, gradients, and FLOPs."""
+        return 0.0, 0.0, 0.0, 0.0  # return (num_layers, num_params, num_gradients, num_flops)
+
+    @staticmethod
+    def iterative_sigma_clipping(data: np.ndarray, sigma: float = 2, max_iters: int = 3):
+        """Apply iterative sigma clipping to data to remove outliers.
+
+        Args:
+            data (np.ndarray): Input data array.
+            sigma (float): Number of standard deviations to use for clipping.
+            max_iters (int): Maximum number of iterations for the clipping process.
+
+        Returns:
+            (np.ndarray): Clipped data array with outliers removed.
+        """
+        data = np.array(data)
+        for _ in range(max_iters):
+            mean, std = np.mean(data), np.std(data)
+            clipped_data = data[(data > mean - sigma * std) & (data < mean + sigma * std)]
+            if len(clipped_data) == len(data):
+                break
+            data = clipped_data
+        return data
+
+    def profile_tensorrt_model(self, engine_file: str, eps: float = 1e-3):
+        """Profile YOLO model performance with TensorRT, measuring average run time and standard deviation.
+
+        Args:
+            engine_file (str): Path to the TensorRT engine file.
+            eps (float): Small epsilon value to prevent division by zero.
+
+        Returns:
+            (tuple[float, float]): Mean and standard deviation of inference time in milliseconds.
+        """
+        if not self.trt or not Path(engine_file).is_file():
+            return 0.0, 0.0
+
+        # Model and input
+        model = YOLO(engine_file)
+        input_data = np.zeros((self.imgsz, self.imgsz, 3), dtype=np.uint8)  # use uint8 for Classify
+
+        # Warmup runs
+        elapsed = 0.0
+        for _ in range(3):
+            start_time = time.time()
+            for _ in range(self.num_warmup_runs):
+                model(input_data, imgsz=self.imgsz, verbose=False)
+            elapsed = time.time() - start_time
+
+        # Compute number of runs as higher of min_time or num_timed_runs
+        num_runs = max(round(self.min_time / (elapsed + eps) * self.num_warmup_runs), self.num_timed_runs * 50)
+
+        # Timed runs
+        run_times = []
+        for _ in TQDM(range(num_runs), desc=engine_file):
+            results = model(input_data, imgsz=self.imgsz, verbose=False)
+            run_times.append(results[0].speed["inference"])  # Convert to milliseconds
+
+        run_times = self.iterative_sigma_clipping(np.array(run_times), sigma=2, max_iters=3)  # sigma clipping
+        return np.mean(run_times), np.std(run_times)
+
+    @staticmethod
+    def check_dynamic(tensor_shape):
+        """Check whether the tensor shape in the ONNX model is dynamic."""
+        return not all(isinstance(dim, int) and dim >= 0 for dim in tensor_shape)
+
+    def profile_onnx_model(self, onnx_file: str, eps: float = 1e-3):
+        """Profile an ONNX model, measuring average inference time and standard deviation across multiple runs.
+
+        Args:
+            onnx_file (str): Path to the ONNX model file.
+            eps (float): Small epsilon value to prevent division by zero.
+
+        Returns:
+            (tuple[float, float]): Mean and standard deviation of inference time in milliseconds.
+        """
+        check_requirements([("onnxruntime", "onnxruntime-gpu")])  # either package meets requirements
+        import onnxruntime as ort
+
+        # Session with either 'TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        sess_options.intra_op_num_threads = 8  # Limit the number of threads
+        sess = ort.InferenceSession(onnx_file, sess_options, providers=["CPUExecutionProvider"])
+
+        input_data_dict = {}
+        for input_tensor in sess.get_inputs():
+            input_type = input_tensor.type
+            if self.check_dynamic(input_tensor.shape):
+                if len(input_tensor.shape) != 4 and self.check_dynamic(input_tensor.shape[1:]):
+                    raise ValueError(f"Unsupported dynamic shape {input_tensor.shape} of {input_tensor.name}")
+                input_shape = (
+                    (1, 3, self.imgsz, self.imgsz) if len(input_tensor.shape) == 4 else (1, *input_tensor.shape[1:])
+                )
+            else:
+                input_shape = input_tensor.shape
+
+            # Mapping ONNX datatype to numpy datatype
+            if "float16" in input_type:
+                input_dtype = np.float16
+            elif "float" in input_type:
+                input_dtype = np.float32
+            elif "double" in input_type:
+                input_dtype = np.float64
+            elif "int64" in input_type:
+                input_dtype = np.int64
+            elif "int32" in input_type:
+                input_dtype = np.int32
+            else:
+                raise ValueError(f"Unsupported ONNX datatype {input_type}")
+
+            input_data = np.random.rand(*input_shape).astype(input_dtype)
+            input_name = input_tensor.name
+            input_data_dict[input_name] = input_data
+
+        output_name = sess.get_outputs()[0].name
+
+        # Warmup runs
+        elapsed = 0.0
+        for _ in range(3):
+            start_time = time.time()
+            for _ in range(self.num_warmup_runs):
+                sess.run([output_name], input_data_dict)
+            elapsed = time.time() - start_time
+
+        # Compute number of runs as higher of min_time or num_timed_runs
+        num_runs = max(round(self.min_time / (elapsed + eps) * self.num_warmup_runs), self.num_timed_runs)
+
+        # Timed runs
+        run_times = []
+        for _ in TQDM(range(num_runs), desc=onnx_file):
+            start_time = time.time()
+            sess.run([output_name], input_data_dict)
+            run_times.append((time.time() - start_time) * 1000)  # Convert to milliseconds
+
+        run_times = self.iterative_sigma_clipping(np.array(run_times), sigma=2, max_iters=5)  # sigma clipping
+        return np.mean(run_times), np.std(run_times)
+
+    def generate_table_row(
+        self,
+        model_name: str,
+        t_onnx: tuple[float, float],
+        t_engine: tuple[float, float],
+        model_info: tuple[float, float, float, float],
+    ):
+        """Generate a table row string with model performance metrics.
+
+        Args:
+            model_name (str): Name of the model.
+            t_onnx (tuple): ONNX model inference time statistics (mean, std).
+            t_engine (tuple): TensorRT engine inference time statistics (mean, std).
+            model_info (tuple): Model information (layers, params, gradients, flops).
+
+        Returns:
+            (str): Formatted table row string with model metrics.
+        """
+        _layers, params, _gradients, flops = model_info
+        return (
+            f"| {model_name:18s} | {self.imgsz} | - | {t_onnx[0]:.1f}±{t_onnx[1]:.1f} ms | {t_engine[0]:.1f}±"
+            f"{t_engine[1]:.1f} ms | {params / 1e6:.1f} | {flops:.1f} |"
+        )
+
+    @staticmethod
+    def generate_results_dict(
+        model_name: str,
+        t_onnx: tuple[float, float],
+        t_engine: tuple[float, float],
+        model_info: tuple[float, float, float, float],
+    ):
+        """Generate a dictionary of profiling results.
+
+        Args:
+            model_name (str): Name of the model.
+            t_onnx (tuple): ONNX model inference time statistics (mean, std).
+            t_engine (tuple): TensorRT engine inference time statistics (mean, std).
+            model_info (tuple): Model information (layers, params, gradients, flops).
+
+        Returns:
+            (dict): Dictionary containing profiling results.
+        """
+        _layers, params, _gradients, flops = model_info
+        return {
+            "model/name": model_name,
+            "model/parameters": params,
+            "model/GFLOPs": round(flops, 3),
+            "model/speed_ONNX(ms)": round(t_onnx[0], 3),
+            "model/speed_TensorRT(ms)": round(t_engine[0], 3),
+        }
+
+    @staticmethod
+    def print_table(table_rows: list[str]):
+        """Print a formatted table of model profiling results.
+
+        Args:
+            table_rows (list[str]): List of formatted table row strings.
+        """
+        gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "GPU"
+        headers = [
+            "Model",
+            "size<br><sup>(pixels)",
+            "mAP<sup>val<br>50-95",
+            f"Speed<br><sup>CPU ({get_cpu_info()}) ONNX<br>(ms)",
+            f"Speed<br><sup>{gpu} TensorRT<br>(ms)",
+            "params<br><sup>(M)",
+            "FLOPs<br><sup>(B)",
+        ]
+        header = "|" + "|".join(f" {h} " for h in headers) + "|"
+        separator = "|" + "|".join("-" * (len(h) + 2) for h in headers) + "|"
+
+        LOGGER.info(f"\n\n{header}")
+        LOGGER.info(separator)
+        for row in table_rows:
+            LOGGER.info(row)

ultralytics/utils/callbacks/__init__.py

@@ -0,0 +1,5 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from .base import add_integration_callbacks, default_callbacks, get_default_callbacks
+
+__all__ = "add_integration_callbacks", "default_callbacks", "get_default_callbacks"

ultralytics/utils/callbacks/base.py

@@ -0,0 +1,233 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""Base callbacks for Ultralytics training, validation, prediction, and export processes."""
+
+from collections import defaultdict
+from copy import deepcopy
+
+# Trainer callbacks ----------------------------------------------------------------------------------------------------
+
+
+def on_pretrain_routine_start(trainer):
+    """Called before the pretraining routine starts."""
+    pass
+
+
+def on_pretrain_routine_end(trainer):
+    """Called after the pretraining routine ends."""
+    pass
+
+
+def on_train_start(trainer):
+    """Called when the training starts."""
+    pass
+
+
+def on_train_epoch_start(trainer):
+    """Called at the start of each training epoch."""
+    pass
+
+
+def on_train_batch_start(trainer):
+    """Called at the start of each training batch."""
+    pass
+
+
+def optimizer_step(trainer):
+    """Called when the optimizer takes a step."""
+    pass
+
+
+def on_before_zero_grad(trainer):
+    """Called before the gradients are set to zero."""
+    pass
+
+
+def on_train_batch_end(trainer):
+    """Called at the end of each training batch."""
+    pass
+
+
+def on_train_epoch_end(trainer):
+    """Called at the end of each training epoch."""
+    pass
+
+
+def on_fit_epoch_end(trainer):
+    """Called at the end of each fit epoch (train + val)."""
+    pass
+
+
+def on_model_save(trainer):
+    """Called when the model is saved."""
+    pass
+
+
+def on_train_end(trainer):
+    """Called when the training ends."""
+    pass
+
+
+def on_params_update(trainer):
+    """Called when the model parameters are updated."""
+    pass
+
+
+def teardown(trainer):
+    """Called during the teardown of the training process."""
+    pass
+
+
+# Validator callbacks --------------------------------------------------------------------------------------------------
+
+
+def on_val_start(validator):
+    """Called when the validation starts."""
+    pass
+
+
+def on_val_batch_start(validator):
+    """Called at the start of each validation batch."""
+    pass
+
+
+def on_val_batch_end(validator):
+    """Called at the end of each validation batch."""
+    pass
+
+
+def on_val_end(validator):
+    """Called when the validation ends."""
+    pass
+
+
+# Predictor callbacks --------------------------------------------------------------------------------------------------
+
+
+def on_predict_start(predictor):
+    """Called when the prediction starts."""
+    pass
+
+
+def on_predict_batch_start(predictor):
+    """Called at the start of each prediction batch."""
+    pass
+
+
+def on_predict_batch_end(predictor):
+    """Called at the end of each prediction batch."""
+    pass
+
+
+def on_predict_postprocess_end(predictor):
+    """Called after the post-processing of the prediction ends."""
+    pass
+
+
+def on_predict_end(predictor):
+    """Called when the prediction ends."""
+    pass
+
+
+# Exporter callbacks ---------------------------------------------------------------------------------------------------
+
+
+def on_export_start(exporter):
+    """Called when the model export starts."""
+    pass
+
+
+def on_export_end(exporter):
+    """Called when the model export ends."""
+    pass
+
+
+default_callbacks = {
+    # Run in trainer
+    "on_pretrain_routine_start": [on_pretrain_routine_start],
+    "on_pretrain_routine_end": [on_pretrain_routine_end],
+    "on_train_start": [on_train_start],
+    "on_train_epoch_start": [on_train_epoch_start],
+    "on_train_batch_start": [on_train_batch_start],
+    "optimizer_step": [optimizer_step],
+    "on_before_zero_grad": [on_before_zero_grad],
+    "on_train_batch_end": [on_train_batch_end],
+    "on_train_epoch_end": [on_train_epoch_end],
+    "on_fit_epoch_end": [on_fit_epoch_end],  # fit = train + val
+    "on_model_save": [on_model_save],
+    "on_train_end": [on_train_end],
+    "on_params_update": [on_params_update],
+    "teardown": [teardown],
+    # Run in validator
+    "on_val_start": [on_val_start],
+    "on_val_batch_start": [on_val_batch_start],
+    "on_val_batch_end": [on_val_batch_end],
+    "on_val_end": [on_val_end],
+    # Run in predictor
+    "on_predict_start": [on_predict_start],
+    "on_predict_batch_start": [on_predict_batch_start],
+    "on_predict_postprocess_end": [on_predict_postprocess_end],
+    "on_predict_batch_end": [on_predict_batch_end],
+    "on_predict_end": [on_predict_end],
+    # Run in exporter
+    "on_export_start": [on_export_start],
+    "on_export_end": [on_export_end],
+}
+
+
+def get_default_callbacks():
+    """Get the default callbacks for Ultralytics training, validation, prediction, and export processes.
+
+    Returns:
+        (dict): Dictionary of default callbacks for various training events. Each key represents an event during the
+            training process, and the corresponding value is a list of callback functions executed when that
+            event occurs.
+
+    Examples:
+        >>> callbacks = get_default_callbacks()
+        >>> print(list(callbacks.keys()))  # show all available callback events
+        ['on_pretrain_routine_start', 'on_pretrain_routine_end', ...]
+    """
+    return defaultdict(list, deepcopy(default_callbacks))
+
+
+def add_integration_callbacks(instance):
+    """Add integration callbacks to the instance's callbacks dictionary.
+
+    This function loads and adds various integration callbacks to the provided instance. The specific callbacks added
+    depend on the type of instance provided. All instances receive HUB callbacks, while Trainer instances also receive
+    additional callbacks for various integrations like ClearML, Comet, DVC, MLflow, Neptune, Ray Tune, TensorBoard, and
+    Weights & Biases.
+
+    Args:
+        instance (Trainer | Predictor | Validator | Exporter): The object instance to which callbacks will be added. The
+            type of instance determines which callbacks are loaded.
+
+    Examples:
+        >>> from ultralytics.engine.trainer import BaseTrainer
+        >>> trainer = BaseTrainer()
+        >>> add_integration_callbacks(trainer)
+    """
+    from .hub import callbacks as hub_cb
+    from .platform import callbacks as platform_cb
+
+    # Load Ultralytics callbacks
+    callbacks_list = [hub_cb, platform_cb]
+
+    # Load training callbacks
+    if "Trainer" in instance.__class__.__name__:
+        from .clearml import callbacks as clear_cb
+        from .comet import callbacks as comet_cb
+        from .dvc import callbacks as dvc_cb
+        from .mlflow import callbacks as mlflow_cb
+        from .neptune import callbacks as neptune_cb
+        from .raytune import callbacks as tune_cb
+        from .tensorboard import callbacks as tb_cb
+        from .wb import callbacks as wb_cb
+
+        callbacks_list.extend([clear_cb, comet_cb, dvc_cb, mlflow_cb, neptune_cb, tune_cb, tb_cb, wb_cb])
+
+    # Add the callbacks to the callbacks dictionary
+    for callbacks in callbacks_list:
+        for k, v in callbacks.items():
+            if v not in instance.callbacks[k]:
+                instance.callbacks[k].append(v)

ultralytics/utils/callbacks/clearml.py

@@ -0,0 +1,146 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from ultralytics.utils import LOGGER, SETTINGS, TESTS_RUNNING
+
+try:
+    assert not TESTS_RUNNING  # do not log pytest
+    assert SETTINGS["clearml"] is True  # verify integration is enabled
+    import clearml
+    from clearml import Task
+
+    assert hasattr(clearml, "__version__")  # verify package is not directory
+
+except (ImportError, AssertionError):
+    clearml = None
+
+
+def _log_debug_samples(files, title: str = "Debug Samples") -> None:
+    """Log files (images) as debug samples in the ClearML task.
+
+    Args:
+        files (list[Path]): A list of file paths in PosixPath format.
+        title (str): A title that groups together images with the same values.
+    """
+    import re
+
+    if task := Task.current_task():
+        for f in files:
+            if f.exists():
+                it = re.search(r"_batch(\d+)", f.name)
+                iteration = int(it.groups()[0]) if it else 0
+                task.get_logger().report_image(
+                    title=title, series=f.name.replace(it.group(), ""), local_path=str(f), iteration=iteration
+                )
+
+
+def _log_plot(title: str, plot_path: str) -> None:
+    """Log an image as a plot in the plot section of ClearML.
+
+    Args:
+        title (str): The title of the plot.
+        plot_path (str | Path): The path to the saved image file.
+    """
+    import matplotlib.image as mpimg
+    import matplotlib.pyplot as plt
+
+    img = mpimg.imread(plot_path)
+    fig = plt.figure()
+    ax = fig.add_axes([0, 0, 1, 1], frameon=False, aspect="auto", xticks=[], yticks=[])  # no ticks
+    ax.imshow(img)
+
+    Task.current_task().get_logger().report_matplotlib_figure(
+        title=title, series="", figure=fig, report_interactive=False
+    )
+
+
+def on_pretrain_routine_start(trainer) -> None:
+    """Initialize and connect ClearML task at the start of pretraining routine."""
+    try:
+        if task := Task.current_task():
+            # WARNING: make sure the automatic pytorch and matplotlib bindings are disabled!
+            # We are logging these plots and model files manually in the integration
+            from clearml.binding.frameworks.pytorch_bind import PatchPyTorchModelIO
+            from clearml.binding.matplotlib_bind import PatchedMatplotlib
+
+            PatchPyTorchModelIO.update_current_task(None)
+            PatchedMatplotlib.update_current_task(None)
+        else:
+            task = Task.init(
+                project_name=trainer.args.project or "Ultralytics",
+                task_name=trainer.args.name,
+                tags=["Ultralytics"],
+                output_uri=True,
+                reuse_last_task_id=False,
+                auto_connect_frameworks={"pytorch": False, "matplotlib": False},
+            )
+            LOGGER.warning(
+                "ClearML Initialized a new task. If you want to run remotely, "
+                "please add clearml-init and connect your arguments before initializing YOLO."
+            )
+        task.connect(vars(trainer.args), name="General")
+    except Exception as e:
+        LOGGER.warning(f"ClearML installed but not initialized correctly, not logging this run. {e}")
+
+
+def on_train_epoch_end(trainer) -> None:
+    """Log debug samples for the first epoch and report current training progress."""
+    if task := Task.current_task():
+        # Log debug samples for first epoch only
+        if trainer.epoch == 1:
+            _log_debug_samples(sorted(trainer.save_dir.glob("train_batch*.jpg")), "Mosaic")
+        # Report the current training progress
+        for k, v in trainer.label_loss_items(trainer.tloss, prefix="train").items():
+            task.get_logger().report_scalar("train", k, v, iteration=trainer.epoch)
+        for k, v in trainer.lr.items():
+            task.get_logger().report_scalar("lr", k, v, iteration=trainer.epoch)
+
+
+def on_fit_epoch_end(trainer) -> None:
+    """Report model information and metrics to logger at the end of an epoch."""
+    if task := Task.current_task():
+        # Report epoch time and validation metrics
+        task.get_logger().report_scalar(
+            title="Epoch Time", series="Epoch Time", value=trainer.epoch_time, iteration=trainer.epoch
+        )
+        for k, v in trainer.metrics.items():
+            title = k.split("/")[0]
+            task.get_logger().report_scalar(title, k, v, iteration=trainer.epoch)
+        if trainer.epoch == 0:
+            from ultralytics.utils.torch_utils import model_info_for_loggers
+
+            for k, v in model_info_for_loggers(trainer).items():
+                task.get_logger().report_single_value(k, v)
+
+
+def on_val_end(validator) -> None:
+    """Log validation results including labels and predictions."""
+    if Task.current_task():
+        # Log validation labels and predictions
+        _log_debug_samples(sorted(validator.save_dir.glob("val*.jpg")), "Validation")
+
+
+def on_train_end(trainer) -> None:
+    """Log final model and training results on training completion."""
+    if task := Task.current_task():
+        # Log final results, confusion matrix and PR plots
+        for f in [*trainer.plots.keys(), *trainer.validator.plots.keys()]:
+            if "batch" not in f.name:
+                _log_plot(title=f.stem, plot_path=f)
+        # Report final metrics
+        for k, v in trainer.validator.metrics.results_dict.items():
+            task.get_logger().report_single_value(k, v)
+        # Log the final model
+        task.update_output_model(model_path=str(trainer.best), model_name=trainer.args.name, auto_delete_file=False)
+
+
+callbacks = (
+    {
+        "on_pretrain_routine_start": on_pretrain_routine_start,
+        "on_train_epoch_end": on_train_epoch_end,
+        "on_fit_epoch_end": on_fit_epoch_end,
+        "on_val_end": on_val_end,
+        "on_train_end": on_train_end,
+    }
+    if clearml
+    else {}
+)

ultralytics/utils/callbacks/comet.py

@@ -0,0 +1,622 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from types import SimpleNamespace
+from typing import Any
+
+import cv2
+import numpy as np
+
+from ultralytics.utils import LOGGER, RANK, SETTINGS, TESTS_RUNNING, ops
+from ultralytics.utils.metrics import ClassifyMetrics, DetMetrics, OBBMetrics, PoseMetrics, SegmentMetrics
+
+try:
+    assert not TESTS_RUNNING  # do not log pytest
+    assert SETTINGS["comet"] is True  # verify integration is enabled
+    import comet_ml
+
+    assert hasattr(comet_ml, "__version__")  # verify package is not directory
+
+    import os
+    from pathlib import Path
+
+    # Ensures certain logging functions only run for supported tasks
+    COMET_SUPPORTED_TASKS = ["detect", "segment"]
+
+    # Names of plots created by Ultralytics that are logged to Comet
+    CONFUSION_MATRIX_PLOT_NAMES = "confusion_matrix", "confusion_matrix_normalized"
+    EVALUATION_PLOT_NAMES = "F1_curve", "P_curve", "R_curve", "PR_curve"
+    LABEL_PLOT_NAMES = ["labels"]
+    SEGMENT_METRICS_PLOT_PREFIX = "Box", "Mask"
+    POSE_METRICS_PLOT_PREFIX = "Box", "Pose"
+    DETECTION_METRICS_PLOT_PREFIX = ["Box"]
+    RESULTS_TABLE_NAME = "results.csv"
+    ARGS_YAML_NAME = "args.yaml"
+
+    _comet_image_prediction_count = 0
+
+except (ImportError, AssertionError):
+    comet_ml = None
+
+
+def _get_comet_mode() -> str:
+    """Return the Comet mode from environment variables, defaulting to 'online'."""
+    comet_mode = os.getenv("COMET_MODE")
+    if comet_mode is not None:
+        LOGGER.warning(
+            "The COMET_MODE environment variable is deprecated. "
+            "Please use COMET_START_ONLINE to set the Comet experiment mode. "
+            "To start an offline Comet experiment, use 'export COMET_START_ONLINE=0'. "
+            "If COMET_START_ONLINE is not set or is set to '1', an online Comet experiment will be created."
+        )
+        return comet_mode
+
+    return "online"
+
+
+def _get_comet_model_name() -> str:
+    """Return the Comet model name from environment variable or default to 'Ultralytics'."""
+    return os.getenv("COMET_MODEL_NAME", "Ultralytics")
+
+
+def _get_eval_batch_logging_interval() -> int:
+    """Get the evaluation batch logging interval from environment variable or use default value 1."""
+    return int(os.getenv("COMET_EVAL_BATCH_LOGGING_INTERVAL", 1))
+
+
+def _get_max_image_predictions_to_log() -> int:
+    """Get the maximum number of image predictions to log from environment variables."""
+    return int(os.getenv("COMET_MAX_IMAGE_PREDICTIONS", 100))
+
+
+def _scale_confidence_score(score: float) -> float:
+    """Scale the confidence score by a factor specified in environment variable."""
+    scale = float(os.getenv("COMET_MAX_CONFIDENCE_SCORE", 100.0))
+    return score * scale
+
+
+def _should_log_confusion_matrix() -> bool:
+    """Determine if the confusion matrix should be logged based on environment variable settings."""
+    return os.getenv("COMET_EVAL_LOG_CONFUSION_MATRIX", "false").lower() == "true"
+
+
+def _should_log_image_predictions() -> bool:
+    """Determine whether to log image predictions based on environment variable."""
+    return os.getenv("COMET_EVAL_LOG_IMAGE_PREDICTIONS", "true").lower() == "true"
+
+
+def _resume_or_create_experiment(args: SimpleNamespace) -> None:
+    """Resume CometML experiment or create a new experiment based on args.
+
+    Ensures that the experiment object is only created in a single process during distributed training.
+
+    Args:
+        args (SimpleNamespace): Training arguments containing project configuration and other parameters.
+    """
+    if RANK not in {-1, 0}:
+        return
+
+    # Set environment variable (if not set by the user) to configure the Comet experiment's online mode under the hood.
+    # IF COMET_START_ONLINE is set by the user it will override COMET_MODE value.
+    if os.getenv("COMET_START_ONLINE") is None:
+        comet_mode = _get_comet_mode()
+        os.environ["COMET_START_ONLINE"] = "1" if comet_mode != "offline" else "0"
+
+    try:
+        _project_name = os.getenv("COMET_PROJECT_NAME", args.project)
+        experiment = comet_ml.start(project_name=_project_name)
+        experiment.log_parameters(vars(args))
+        experiment.log_others(
+            {
+                "eval_batch_logging_interval": _get_eval_batch_logging_interval(),
+                "log_confusion_matrix_on_eval": _should_log_confusion_matrix(),
+                "log_image_predictions": _should_log_image_predictions(),
+                "max_image_predictions": _get_max_image_predictions_to_log(),
+            }
+        )
+        experiment.log_other("Created from", "ultralytics")
+
+    except Exception as e:
+        LOGGER.warning(f"Comet installed but not initialized correctly, not logging this run. {e}")
+
+
+def _fetch_trainer_metadata(trainer) -> dict:
+    """Return metadata for YOLO training including epoch and asset saving status.
+
+    Args:
+        trainer (ultralytics.engine.trainer.BaseTrainer): The YOLO trainer object containing training state and config.
+
+    Returns:
+        (dict): Dictionary containing current epoch, step, save assets flag, and final epoch flag.
+    """
+    curr_epoch = trainer.epoch + 1
+
+    train_num_steps_per_epoch = len(trainer.train_loader.dataset) // trainer.batch_size
+    curr_step = curr_epoch * train_num_steps_per_epoch
+    final_epoch = curr_epoch == trainer.epochs
+
+    save = trainer.args.save
+    save_period = trainer.args.save_period
+    save_interval = curr_epoch % save_period == 0
+    save_assets = save and save_period > 0 and save_interval and not final_epoch
+
+    return dict(curr_epoch=curr_epoch, curr_step=curr_step, save_assets=save_assets, final_epoch=final_epoch)
+
+
+def _scale_bounding_box_to_original_image_shape(
+    box, resized_image_shape, original_image_shape, ratio_pad
+) -> list[float]:
+    """Scale bounding box from resized image coordinates to original image coordinates.
+
+    YOLO resizes images during training and the label values are normalized based on this resized shape. This function
+    rescales the bounding box labels to the original image shape.
+
+    Args:
+        box (torch.Tensor): Bounding box in normalized xywh format.
+        resized_image_shape (tuple): Shape of the resized image (height, width).
+        original_image_shape (tuple): Shape of the original image (height, width).
+        ratio_pad (tuple): Ratio and padding information for scaling.
+
+    Returns:
+        (list[float]): Scaled bounding box coordinates in xywh format with top-left corner adjustment.
+    """
+    resized_image_height, resized_image_width = resized_image_shape
+
+    # Convert normalized xywh format predictions to xyxy in resized scale format
+    box = ops.xywhn2xyxy(box, h=resized_image_height, w=resized_image_width)
+    # Scale box predictions from resized image scale back to original image scale
+    box = ops.scale_boxes(resized_image_shape, box, original_image_shape, ratio_pad)
+    # Convert bounding box format from xyxy to xywh for Comet logging
+    box = ops.xyxy2xywh(box)
+    # Adjust xy center to correspond top-left corner
+    box[:2] -= box[2:] / 2
+    box = box.tolist()
+
+    return box
+
+
+def _format_ground_truth_annotations_for_detection(img_idx, image_path, batch, class_name_map=None) -> dict | None:
+    """Format ground truth annotations for object detection.
+
+    This function processes ground truth annotations from a batch of images for object detection tasks. It extracts
+    bounding boxes, class labels, and other metadata for a specific image in the batch, and formats them for
+    visualization or evaluation.
+
+    Args:
+        img_idx (int): Index of the image in the batch to process.
+        image_path (str | Path): Path to the image file.
+        batch (dict): Batch dictionary containing detection data with keys:
+            - 'batch_idx': Tensor of batch indices
+            - 'bboxes': Tensor of bounding boxes in normalized xywh format
+            - 'cls': Tensor of class labels
+            - 'ori_shape': Original image shapes
+            - 'resized_shape': Resized image shapes
+            - 'ratio_pad': Ratio and padding information
+        class_name_map (dict, optional): Mapping from class indices to class names.
+
+    Returns:
+        (dict | None): Formatted ground truth annotations with keys 'name' and 'data', where 'data' is a list of
+            annotation dicts each containing 'boxes', 'label', and 'score' keys. Returns None if no bounding boxes are
+            found for the image.
+    """
+    indices = batch["batch_idx"] == img_idx
+    bboxes = batch["bboxes"][indices]
+    if len(bboxes) == 0:
+        LOGGER.debug(f"Comet Image: {image_path} has no bounding boxes labels")
+        return None
+
+    cls_labels = batch["cls"][indices].squeeze(1).tolist()
+    if class_name_map:
+        cls_labels = [str(class_name_map[label]) for label in cls_labels]
+
+    original_image_shape = batch["ori_shape"][img_idx]
+    resized_image_shape = batch["resized_shape"][img_idx]
+    ratio_pad = batch["ratio_pad"][img_idx]
+
+    data = []
+    for box, label in zip(bboxes, cls_labels):
+        box = _scale_bounding_box_to_original_image_shape(box, resized_image_shape, original_image_shape, ratio_pad)
+        data.append(
+            {
+                "boxes": [box],
+                "label": f"gt_{label}",
+                "score": _scale_confidence_score(1.0),
+            }
+        )
+
+    return {"name": "ground_truth", "data": data}
+
+
+def _format_prediction_annotations(image_path, metadata, class_label_map=None, class_map=None) -> dict | None:
+    """Format YOLO predictions for object detection visualization.
+
+    Args:
+        image_path (Path): Path to the image file.
+        metadata (dict): Prediction metadata containing bounding boxes and class information.
+        class_label_map (dict, optional): Mapping from class indices to class names.
+        class_map (dict, optional): Additional class mapping for label conversion.
+
+    Returns:
+        (dict | None): Formatted prediction annotations or None if no predictions exist.
+    """
+    stem = image_path.stem
+    image_id = int(stem) if stem.isnumeric() else stem
+
+    predictions = metadata.get(image_id)
+    if not predictions:
+        LOGGER.debug(f"Comet Image: {image_path} has no bounding boxes predictions")
+        return None
+
+    # apply the mapping that was used to map the predicted classes when the JSON was created
+    if class_label_map and class_map:
+        class_label_map = {class_map[k]: v for k, v in class_label_map.items()}
+    try:
+        # import pycotools utilities to decompress annotations for various tasks, e.g. segmentation
+        from faster_coco_eval.core.mask import decode
+    except ImportError:
+        decode = None
+
+    data = []
+    for prediction in predictions:
+        boxes = prediction["bbox"]
+        score = _scale_confidence_score(prediction["score"])
+        cls_label = prediction["category_id"]
+        if class_label_map:
+            cls_label = str(class_label_map[cls_label])
+
+        annotation_data = {"boxes": [boxes], "label": cls_label, "score": score}
+
+        if decode is not None:
+            # do segmentation processing only if we are able to decode it
+            segments = prediction.get("segmentation", None)
+            if segments is not None:
+                segments = _extract_segmentation_annotation(segments, decode)
+            if segments is not None:
+                annotation_data["points"] = segments
+
+        data.append(annotation_data)
+
+    return {"name": "prediction", "data": data}
+
+
+def _extract_segmentation_annotation(segmentation_raw: str, decode: Callable) -> list[list[Any]] | None:
+    """Extract segmentation annotation from compressed segmentations as list of polygons.
+
+    Args:
+        segmentation_raw (str): Raw segmentation data in compressed format.
+        decode (Callable): Function to decode the compressed segmentation data.
+
+    Returns:
+        (list[list[Any]] | None): List of polygon points or None if extraction fails.
+    """
+    try:
+        mask = decode(segmentation_raw)
+        contours, _ = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+        annotations = [np.array(polygon).squeeze() for polygon in contours if len(polygon) >= 3]
+        return [annotation.ravel().tolist() for annotation in annotations]
+    except Exception as e:
+        LOGGER.warning(f"Comet Failed to extract segmentation annotation: {e}")
+    return None
+
+
+def _fetch_annotations(img_idx, image_path, batch, prediction_metadata_map, class_label_map, class_map) -> list | None:
+    """Join the ground truth and prediction annotations if they exist.
+
+    Args:
+        img_idx (int): Index of the image in the batch.
+        image_path (Path): Path to the image file.
+        batch (dict): Batch data containing ground truth annotations.
+        prediction_metadata_map (dict): Map of prediction metadata by image ID.
+        class_label_map (dict): Mapping from class indices to class names.
+        class_map (dict): Additional class mapping for label conversion.
+
+    Returns:
+        (list | None): List of annotation dictionaries or None if no annotations exist.
+    """
+    ground_truth_annotations = _format_ground_truth_annotations_for_detection(
+        img_idx, image_path, batch, class_label_map
+    )
+    prediction_annotations = _format_prediction_annotations(
+        image_path, prediction_metadata_map, class_label_map, class_map
+    )
+
+    annotations = [
+        annotation for annotation in [ground_truth_annotations, prediction_annotations] if annotation is not None
+    ]
+    return [annotations] if annotations else None
+
+
+def _create_prediction_metadata_map(model_predictions) -> dict:
+    """Create metadata map for model predictions by grouping them based on image ID."""
+    pred_metadata_map = {}
+    for prediction in model_predictions:
+        pred_metadata_map.setdefault(prediction["image_id"], [])
+        pred_metadata_map[prediction["image_id"]].append(prediction)
+
+    return pred_metadata_map
+
+
+def _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch) -> None:
+    """Log the confusion matrix to Comet experiment."""
+    conf_mat = trainer.validator.confusion_matrix.matrix
+    names = [*list(trainer.data["names"].values()), "background"]
+    experiment.log_confusion_matrix(
+        matrix=conf_mat, labels=names, max_categories=len(names), epoch=curr_epoch, step=curr_step
+    )
+
+
+def _log_images(experiment, image_paths, curr_step: int | None, annotations=None) -> None:
+    """Log images to the experiment with optional annotations.
+
+    This function logs images to a Comet ML experiment, optionally including annotation data for visualization such as
+    bounding boxes or segmentation masks.
+
+    Args:
+        experiment (comet_ml.CometExperiment): The Comet ML experiment to log images to.
+        image_paths (list[Path]): List of paths to images that will be logged.
+        curr_step (int | None): Current training step/iteration for tracking in the experiment timeline.
+        annotations (list[list[dict]], optional): Nested list of annotation dictionaries for each image. Each annotation
+            contains visualization data like bounding boxes, labels, and confidence scores.
+    """
+    if annotations:
+        for image_path, annotation in zip(image_paths, annotations):
+            experiment.log_image(image_path, name=image_path.stem, step=curr_step, annotations=annotation)
+
+    else:
+        for image_path in image_paths:
+            experiment.log_image(image_path, name=image_path.stem, step=curr_step)
+
+
+def _log_image_predictions(experiment, validator, curr_step) -> None:
+    """Log image predictions to a Comet ML experiment during model validation.
+
+    This function processes validation data and formats both ground truth and prediction annotations for visualization
+    in the Comet dashboard. The function respects configured limits on the number of images to log.
+
+    Args:
+        experiment (comet_ml.CometExperiment): The Comet ML experiment to log to.
+        validator (BaseValidator): The validator instance containing validation data and predictions.
+        curr_step (int): The current training step for logging timeline.
+
+    Notes:
+        This function uses global state to track the number of logged predictions across calls.
+        It only logs predictions for supported tasks defined in COMET_SUPPORTED_TASKS.
+        The number of logged images is limited by the COMET_MAX_IMAGE_PREDICTIONS environment variable.
+    """
+    global _comet_image_prediction_count
+
+    task = validator.args.task
+    if task not in COMET_SUPPORTED_TASKS:
+        return
+
+    jdict = validator.jdict
+    if not jdict:
+        return
+
+    predictions_metadata_map = _create_prediction_metadata_map(jdict)
+    dataloader = validator.dataloader
+    class_label_map = validator.names
+    class_map = getattr(validator, "class_map", None)
+
+    batch_logging_interval = _get_eval_batch_logging_interval()
+    max_image_predictions = _get_max_image_predictions_to_log()
+
+    for batch_idx, batch in enumerate(dataloader):
+        if (batch_idx + 1) % batch_logging_interval != 0:
+            continue
+
+        image_paths = batch["im_file"]
+        for img_idx, image_path in enumerate(image_paths):
+            if _comet_image_prediction_count >= max_image_predictions:
+                return
+
+            image_path = Path(image_path)
+            annotations = _fetch_annotations(
+                img_idx,
+                image_path,
+                batch,
+                predictions_metadata_map,
+                class_label_map,
+                class_map=class_map,
+            )
+            _log_images(
+                experiment,
+                [image_path],
+                curr_step,
+                annotations=annotations,
+            )
+            _comet_image_prediction_count += 1
+
+
+def _log_plots(experiment, trainer) -> None:
+    """Log evaluation plots and label plots for the experiment.
+
+    This function logs various evaluation plots and confusion matrices to the experiment tracking system. It handles
+    different types of metrics (SegmentMetrics, PoseMetrics, DetMetrics, OBBMetrics) and logs the appropriate plots for
+    each type.
+
+    Args:
+        experiment (comet_ml.CometExperiment): The Comet ML experiment to log plots to.
+        trainer (ultralytics.engine.trainer.BaseTrainer): The trainer object containing validation metrics and save
+            directory information.
+
+    Examples:
+        >>> from ultralytics.utils.callbacks.comet import _log_plots
+        >>> _log_plots(experiment, trainer)
+    """
+    plot_filenames = None
+    if isinstance(trainer.validator.metrics, SegmentMetrics):
+        plot_filenames = [
+            trainer.save_dir / f"{prefix}{plots}.png"
+            for plots in EVALUATION_PLOT_NAMES
+            for prefix in SEGMENT_METRICS_PLOT_PREFIX
+        ]
+    elif isinstance(trainer.validator.metrics, PoseMetrics):
+        plot_filenames = [
+            trainer.save_dir / f"{prefix}{plots}.png"
+            for plots in EVALUATION_PLOT_NAMES
+            for prefix in POSE_METRICS_PLOT_PREFIX
+        ]
+    elif isinstance(trainer.validator.metrics, (DetMetrics, OBBMetrics)):
+        plot_filenames = [
+            trainer.save_dir / f"{prefix}{plots}.png"
+            for plots in EVALUATION_PLOT_NAMES
+            for prefix in DETECTION_METRICS_PLOT_PREFIX
+        ]
+
+    if plot_filenames is not None:
+        _log_images(experiment, plot_filenames, None)
+
+    confusion_matrix_filenames = [trainer.save_dir / f"{plots}.png" for plots in CONFUSION_MATRIX_PLOT_NAMES]
+    _log_images(experiment, confusion_matrix_filenames, None)
+
+    if not isinstance(trainer.validator.metrics, ClassifyMetrics):
+        label_plot_filenames = [trainer.save_dir / f"{labels}.jpg" for labels in LABEL_PLOT_NAMES]
+        _log_images(experiment, label_plot_filenames, None)
+
+
+def _log_model(experiment, trainer) -> None:
+    """Log the best-trained model to Comet.ml."""
+    model_name = _get_comet_model_name()
+    experiment.log_model(model_name, file_or_folder=str(trainer.best), file_name="best.pt", overwrite=True)
+
+
+def _log_image_batches(experiment, trainer, curr_step: int) -> None:
+    """Log samples of image batches for train and validation."""
+    _log_images(experiment, trainer.save_dir.glob("train_batch*.jpg"), curr_step)
+    _log_images(experiment, trainer.save_dir.glob("val_batch*.jpg"), curr_step)
+
+
+def _log_asset(experiment, asset_path) -> None:
+    """Logs a specific asset file to the given experiment.
+
+    This function facilitates logging an asset, such as a file, to the provided
+    experiment. It enables integration with experiment tracking platforms.
+
+    Args:
+        experiment (comet_ml.CometExperiment): The experiment instance to which the asset will be logged.
+        asset_path (Path): The file path of the asset to log.
+    """
+    experiment.log_asset(asset_path)
+
+
+def _log_table(experiment, table_path) -> None:
+    """Logs a table to the provided experiment.
+
+    This function is used to log a table file to the given experiment. The table is identified by its file path.
+
+    Args:
+        experiment (comet_ml.CometExperiment): The experiment object where the table file will be logged.
+        table_path (Path): The file path of the table to be logged.
+    """
+    experiment.log_table(str(table_path))
+
+
+def on_pretrain_routine_start(trainer) -> None:
+    """Create or resume a CometML experiment at the start of a YOLO pre-training routine."""
+    _resume_or_create_experiment(trainer.args)
+
+
+def on_train_epoch_end(trainer) -> None:
+    """Log metrics and save batch images at the end of training epochs."""
+    experiment = comet_ml.get_running_experiment()
+    if not experiment:
+        return
+
+    metadata = _fetch_trainer_metadata(trainer)
+    curr_epoch = metadata["curr_epoch"]
+    curr_step = metadata["curr_step"]
+
+    experiment.log_metrics(trainer.label_loss_items(trainer.tloss, prefix="train"), step=curr_step, epoch=curr_epoch)
+
+
+def on_fit_epoch_end(trainer) -> None:
+    """Log model assets at the end of each epoch during training.
+
+    This function is called at the end of each training epoch to log metrics, learning rates, and model information to a
+    Comet ML experiment. It also logs model assets, confusion matrices, and image predictions based on configuration
+    settings.
+
+    The function retrieves the current Comet ML experiment and logs various training metrics. If it's the first epoch,
+    it also logs model information. On specified save intervals, it logs the model, confusion matrix (if enabled), and
+    image predictions (if enabled).
+
+    Args:
+        trainer (BaseTrainer): The YOLO trainer object containing training state, metrics, and configuration.
+
+    Examples:
+        >>> # Inside a training loop
+        >>> on_fit_epoch_end(trainer)  # Log metrics and assets to Comet ML
+    """
+    experiment = comet_ml.get_running_experiment()
+    if not experiment:
+        return
+
+    metadata = _fetch_trainer_metadata(trainer)
+    curr_epoch = metadata["curr_epoch"]
+    curr_step = metadata["curr_step"]
+    save_assets = metadata["save_assets"]
+
+    experiment.log_metrics(trainer.metrics, step=curr_step, epoch=curr_epoch)
+    experiment.log_metrics(trainer.lr, step=curr_step, epoch=curr_epoch)
+    if curr_epoch == 1:
+        from ultralytics.utils.torch_utils import model_info_for_loggers
+
+        experiment.log_metrics(model_info_for_loggers(trainer), step=curr_step, epoch=curr_epoch)
+
+    if not save_assets:
+        return
+
+    _log_model(experiment, trainer)
+    if _should_log_confusion_matrix():
+        _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch)
+    if _should_log_image_predictions():
+        _log_image_predictions(experiment, trainer.validator, curr_step)
+
+
+def on_train_end(trainer) -> None:
+    """Perform operations at the end of training."""
+    experiment = comet_ml.get_running_experiment()
+    if not experiment:
+        return
+
+    metadata = _fetch_trainer_metadata(trainer)
+    curr_epoch = metadata["curr_epoch"]
+    curr_step = metadata["curr_step"]
+    plots = trainer.args.plots
+
+    _log_model(experiment, trainer)
+    if plots:
+        _log_plots(experiment, trainer)
+
+    _log_confusion_matrix(experiment, trainer, curr_step, curr_epoch)
+    _log_image_predictions(experiment, trainer.validator, curr_step)
+    _log_image_batches(experiment, trainer, curr_step)
+    # log results table
+    table_path = trainer.save_dir / RESULTS_TABLE_NAME
+    if table_path.exists():
+        _log_table(experiment, table_path)
+
+    # log arguments YAML
+    args_path = trainer.save_dir / ARGS_YAML_NAME
+    if args_path.exists():
+        _log_asset(experiment, args_path)
+
+    experiment.end()
+
+    global _comet_image_prediction_count
+    _comet_image_prediction_count = 0
+
+
+callbacks = (
+    {
+        "on_pretrain_routine_start": on_pretrain_routine_start,
+        "on_train_epoch_end": on_train_epoch_end,
+        "on_fit_epoch_end": on_fit_epoch_end,
+        "on_train_end": on_train_end,
+    }
+    if comet_ml
+    else {}
+)

ultralytics/utils/callbacks/dvc.py

@@ -0,0 +1,197 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from pathlib import Path
+
+from ultralytics.utils import LOGGER, SETTINGS, TESTS_RUNNING, checks
+
+try:
+    assert not TESTS_RUNNING  # do not log pytest
+    assert SETTINGS["dvc"] is True  # verify integration is enabled
+    import dvclive
+
+    assert checks.check_version("dvclive", "2.11.0", verbose=True)
+
+    import os
+    import re
+
+    # DVCLive logger instance
+    live = None
+    _processed_plots = {}
+
+    # `on_fit_epoch_end` is called on final validation (probably need to be fixed) for now this is the way we
+    # distinguish final evaluation of the best model vs last epoch validation
+    _training_epoch = False
+
+except (ImportError, AssertionError, TypeError):
+    dvclive = None
+
+
+def _log_images(path: Path, prefix: str = "") -> None:
+    """Log images at specified path with an optional prefix using DVCLive.
+
+    This function logs images found at the given path to DVCLive, organizing them by batch to enable slider
+    functionality in the UI. It processes image filenames to extract batch information and restructures the path
+    accordingly.
+
+    Args:
+        path (Path): Path to the image file to be logged.
+        prefix (str, optional): Optional prefix to add to the image name when logging.
+
+    Examples:
+        >>> from pathlib import Path
+        >>> _log_images(Path("runs/train/exp/val_batch0_pred.jpg"), prefix="validation")
+    """
+    if live:
+        name = path.name
+
+        # Group images by batch to enable sliders in UI
+        if m := re.search(r"_batch(\d+)", name):
+            ni = m[1]
+            new_stem = re.sub(r"_batch(\d+)", "_batch", path.stem)
+            name = (Path(new_stem) / ni).with_suffix(path.suffix)
+
+        live.log_image(os.path.join(prefix, name), path)
+
+
+def _log_plots(plots: dict, prefix: str = "") -> None:
+    """Log plot images for training progress if they have not been previously processed.
+
+    Args:
+        plots (dict): Dictionary containing plot information with timestamps.
+        prefix (str, optional): Optional prefix to add to the logged image paths.
+    """
+    for name, params in plots.items():
+        timestamp = params["timestamp"]
+        if _processed_plots.get(name) != timestamp:
+            _log_images(name, prefix)
+            _processed_plots[name] = timestamp
+
+
+def _log_confusion_matrix(validator) -> None:
+    """Log confusion matrix for a validator using DVCLive.
+
+    This function processes the confusion matrix from a validator object and logs it to DVCLive by converting the matrix
+    into lists of target and prediction labels.
+
+    Args:
+        validator (BaseValidator): The validator object containing the confusion matrix and class names. Must have
+            attributes confusion_matrix.matrix, confusion_matrix.task, and names.
+    """
+    targets = []
+    preds = []
+    matrix = validator.confusion_matrix.matrix
+    names = list(validator.names.values())
+    if validator.confusion_matrix.task == "detect":
+        names += ["background"]
+
+    for ti, pred in enumerate(matrix.T.astype(int)):
+        for pi, num in enumerate(pred):
+            targets.extend([names[ti]] * num)
+            preds.extend([names[pi]] * num)
+
+    live.log_sklearn_plot("confusion_matrix", targets, preds, name="cf.json", normalized=True)
+
+
+def on_pretrain_routine_start(trainer) -> None:
+    """Initialize DVCLive logger for training metadata during pre-training routine."""
+    try:
+        global live
+        live = dvclive.Live(save_dvc_exp=True, cache_images=True)
+        LOGGER.info("DVCLive is detected and auto logging is enabled (run 'yolo settings dvc=False' to disable).")
+    except Exception as e:
+        LOGGER.warning(f"DVCLive installed but not initialized correctly, not logging this run. {e}")
+
+
+def on_pretrain_routine_end(trainer) -> None:
+    """Log plots related to the training process at the end of the pretraining routine."""
+    _log_plots(trainer.plots, "train")
+
+
+def on_train_start(trainer) -> None:
+    """Log the training parameters if DVCLive logging is active."""
+    if live:
+        live.log_params(trainer.args)
+
+
+def on_train_epoch_start(trainer) -> None:
+    """Set the global variable _training_epoch value to True at the start of each training epoch."""
+    global _training_epoch
+    _training_epoch = True
+
+
+def on_fit_epoch_end(trainer) -> None:
+    """Log training metrics, model info, and advance to next step at the end of each fit epoch.
+
+    This function is called at the end of each fit epoch during training. It logs various metrics including training
+    loss items, validation metrics, and learning rates. On the first epoch, it also logs model
+    information. Additionally, it logs training and validation plots and advances the DVCLive step counter.
+
+    Args:
+        trainer (BaseTrainer): The trainer object containing training state, metrics, and plots.
+
+    Notes:
+        This function only performs logging operations when DVCLive logging is active and during a training epoch.
+        The global variable _training_epoch is used to track whether the current epoch is a training epoch.
+    """
+    global _training_epoch
+    if live and _training_epoch:
+        all_metrics = {**trainer.label_loss_items(trainer.tloss, prefix="train"), **trainer.metrics, **trainer.lr}
+        for metric, value in all_metrics.items():
+            live.log_metric(metric, value)
+
+        if trainer.epoch == 0:
+            from ultralytics.utils.torch_utils import model_info_for_loggers
+
+            for metric, value in model_info_for_loggers(trainer).items():
+                live.log_metric(metric, value, plot=False)
+
+        _log_plots(trainer.plots, "train")
+        _log_plots(trainer.validator.plots, "val")
+
+        live.next_step()
+        _training_epoch = False
+
+
+def on_train_end(trainer) -> None:
+    """Log best metrics, plots, and confusion matrix at the end of training.
+
+    This function is called at the conclusion of the training process to log final metrics, visualizations, and model
+    artifacts if DVCLive logging is active. It captures the best model performance metrics, training plots, validation
+    plots, and confusion matrix for later analysis.
+
+    Args:
+        trainer (BaseTrainer): The trainer object containing training state, metrics, and validation results.
+
+    Examples:
+        >>> # Inside a custom training loop
+        >>> from ultralytics.utils.callbacks.dvc import on_train_end
+        >>> on_train_end(trainer)  # Log final metrics and artifacts
+    """
+    if live:
+        # At the end log the best metrics. It runs validator on the best model internally.
+        all_metrics = {**trainer.label_loss_items(trainer.tloss, prefix="train"), **trainer.metrics, **trainer.lr}
+        for metric, value in all_metrics.items():
+            live.log_metric(metric, value, plot=False)
+
+        _log_plots(trainer.plots, "val")
+        _log_plots(trainer.validator.plots, "val")
+        _log_confusion_matrix(trainer.validator)
+
+        if trainer.best.exists():
+            live.log_artifact(trainer.best, copy=True, type="model")
+
+        live.end()
+
+
+callbacks = (
+    {
+        "on_pretrain_routine_start": on_pretrain_routine_start,
+        "on_pretrain_routine_end": on_pretrain_routine_end,
+        "on_train_start": on_train_start,
+        "on_train_epoch_start": on_train_epoch_start,
+        "on_fit_epoch_end": on_fit_epoch_end,
+        "on_train_end": on_train_end,
+    }
+    if dvclive
+    else {}
+)

ultralytics/utils/callbacks/hub.py

@@ -0,0 +1,110 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import json
+from time import time
+
+from ultralytics.hub import HUB_WEB_ROOT, PREFIX, HUBTrainingSession
+from ultralytics.utils import LOGGER, RANK, SETTINGS
+from ultralytics.utils.events import events
+
+
+def on_pretrain_routine_start(trainer):
+    """Create a remote Ultralytics HUB session to log local model training."""
+    if RANK in {-1, 0} and SETTINGS["hub"] is True and SETTINGS["api_key"] and trainer.hub_session is None:
+        trainer.hub_session = HUBTrainingSession.create_session(trainer.args.model, trainer.args)
+
+
+def on_pretrain_routine_end(trainer):
+    """Initialize timers for upload rate limiting before training begins."""
+    if session := getattr(trainer, "hub_session", None):
+        # Start timer for upload rate limit
+        session.timers = {"metrics": time(), "ckpt": time()}  # start timer for session rate limiting
+
+
+def on_fit_epoch_end(trainer):
+    """Upload training progress metrics to Ultralytics HUB at the end of each epoch."""
+    if session := getattr(trainer, "hub_session", None):
+        # Upload metrics after validation ends
+        all_plots = {
+            **trainer.label_loss_items(trainer.tloss, prefix="train"),
+            **trainer.metrics,
+        }
+        if trainer.epoch == 0:
+            from ultralytics.utils.torch_utils import model_info_for_loggers
+
+            all_plots = {**all_plots, **model_info_for_loggers(trainer)}
+
+        session.metrics_queue[trainer.epoch] = json.dumps(all_plots)
+
+        # If any metrics failed to upload previously, add them to the queue to attempt uploading again
+        if session.metrics_upload_failed_queue:
+            session.metrics_queue.update(session.metrics_upload_failed_queue)
+
+        if time() - session.timers["metrics"] > session.rate_limits["metrics"]:
+            session.upload_metrics()
+            session.timers["metrics"] = time()  # reset timer
+            session.metrics_queue = {}  # reset queue
+
+
+def on_model_save(trainer):
+    """Upload model checkpoints to Ultralytics HUB with rate limiting."""
+    if session := getattr(trainer, "hub_session", None):
+        # Upload checkpoints with rate limiting
+        is_best = trainer.best_fitness == trainer.fitness
+        if time() - session.timers["ckpt"] > session.rate_limits["ckpt"]:
+            LOGGER.info(f"{PREFIX}Uploading checkpoint {HUB_WEB_ROOT}/models/{session.model.id}")
+            session.upload_model(trainer.epoch, trainer.last, is_best)
+            session.timers["ckpt"] = time()  # reset timer
+
+
+def on_train_end(trainer):
+    """Upload final model and metrics to Ultralytics HUB at the end of training."""
+    if session := getattr(trainer, "hub_session", None):
+        # Upload final model and metrics with exponential standoff
+        LOGGER.info(f"{PREFIX}Syncing final model...")
+        session.upload_model(
+            trainer.epoch,
+            trainer.best,
+            map=trainer.metrics.get("metrics/mAP50-95(B)", 0),
+            final=True,
+        )
+        session.alive = False  # stop heartbeats
+        LOGGER.info(f"{PREFIX}Done ✅\n{PREFIX}View model at {session.model_url} 🚀")
+
+
+def on_train_start(trainer):
+    """Run events on train start."""
+    events(trainer.args, trainer.device)
+
+
+def on_val_start(validator):
+    """Run events on validation start."""
+    if not validator.training:
+        events(validator.args, validator.device)
+
+
+def on_predict_start(predictor):
+    """Run events on predict start."""
+    events(predictor.args, predictor.device)
+
+
+def on_export_start(exporter):
+    """Run events on export start."""
+    events(exporter.args, exporter.device)
+
+
+callbacks = (
+    {
+        "on_pretrain_routine_start": on_pretrain_routine_start,
+        "on_pretrain_routine_end": on_pretrain_routine_end,
+        "on_fit_epoch_end": on_fit_epoch_end,
+        "on_model_save": on_model_save,
+        "on_train_end": on_train_end,
+        "on_train_start": on_train_start,
+        "on_val_start": on_val_start,
+        "on_predict_start": on_predict_start,
+        "on_export_start": on_export_start,
+    }
+    if SETTINGS["hub"] is True
+    else {}
+)

ultralytics/utils/callbacks/mlflow.py

@@ -0,0 +1,134 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""
+MLflow Logging for Ultralytics YOLO.
+
+This module enables MLflow logging for Ultralytics YOLO. It logs metrics, parameters, and model artifacts.
+For setting up, a tracking URI should be specified. The logging can be customized using environment variables.
+
+Commands:
+    1. To set a project name:
+        `export MLFLOW_EXPERIMENT_NAME=<your_experiment_name>` or use the project=<project> argument
+
+    2. To set a run name:
+        `export MLFLOW_RUN=<your_run_name>` or use the name=<name> argument
+
+    3. To start a local MLflow server:
+        mlflow server --backend-store-uri runs/mlflow
+       It will by default start a local server at http://127.0.0.1:5000.
+       To specify a different URI, set the MLFLOW_TRACKING_URI environment variable.
+
+    4. To kill all running MLflow server instances:
+        ps aux | grep 'mlflow' | grep -v 'grep' | awk '{print $2}' | xargs kill -9
+"""
+
+from ultralytics.utils import LOGGER, RUNS_DIR, SETTINGS, TESTS_RUNNING, colorstr
+
+try:
+    import os
+
+    assert not TESTS_RUNNING or "test_mlflow" in os.environ.get("PYTEST_CURRENT_TEST", "")  # do not log pytest
+    assert SETTINGS["mlflow"] is True  # verify integration is enabled
+    import mlflow
+
+    assert hasattr(mlflow, "__version__")  # verify package is not directory
+    from pathlib import Path
+
+    PREFIX = colorstr("MLflow: ")
+
+except (ImportError, AssertionError):
+    mlflow = None
+
+
+def sanitize_dict(x: dict) -> dict:
+    """Sanitize dictionary keys by removing parentheses and converting values to floats."""
+    return {k.replace("(", "").replace(")", ""): float(v) for k, v in x.items()}
+
+
+def on_pretrain_routine_end(trainer):
+    """Log training parameters to MLflow at the end of the pretraining routine.
+
+    This function sets up MLflow logging based on environment variables and trainer arguments. It sets the tracking URI,
+    experiment name, and run name, then starts the MLflow run if not already active. It finally logs the parameters from
+    the trainer.
+
+    Args:
+        trainer (ultralytics.engine.trainer.BaseTrainer): The training object with arguments and parameters to log.
+
+    Notes:
+        MLFLOW_TRACKING_URI: The URI for MLflow tracking. If not set, defaults to 'runs/mlflow'.
+        MLFLOW_EXPERIMENT_NAME: The name of the MLflow experiment. If not set, defaults to trainer.args.project.
+        MLFLOW_RUN: The name of the MLflow run. If not set, defaults to trainer.args.name.
+        MLFLOW_KEEP_RUN_ACTIVE: Boolean indicating whether to keep the MLflow run active after training ends.
+    """
+    global mlflow
+
+    uri = os.environ.get("MLFLOW_TRACKING_URI") or str(RUNS_DIR / "mlflow")
+    LOGGER.debug(f"{PREFIX} tracking uri: {uri}")
+    mlflow.set_tracking_uri(uri)
+
+    # Set experiment and run names
+    experiment_name = os.environ.get("MLFLOW_EXPERIMENT_NAME") or trainer.args.project or "/Shared/Ultralytics"
+    run_name = os.environ.get("MLFLOW_RUN") or trainer.args.name
+    mlflow.set_experiment(experiment_name)
+
+    mlflow.autolog()
+    try:
+        active_run = mlflow.active_run() or mlflow.start_run(run_name=run_name)
+        LOGGER.info(f"{PREFIX}logging run_id({active_run.info.run_id}) to {uri}")
+        if Path(uri).is_dir():
+            LOGGER.info(f"{PREFIX}view at http://127.0.0.1:5000 with 'mlflow server --backend-store-uri {uri}'")
+        LOGGER.info(f"{PREFIX}disable with 'yolo settings mlflow=False'")
+        mlflow.log_params(dict(trainer.args))
+    except Exception as e:
+        LOGGER.warning(f"{PREFIX}Failed to initialize: {e}")
+        LOGGER.warning(f"{PREFIX}Not tracking this run")
+
+
+def on_train_epoch_end(trainer):
+    """Log training metrics at the end of each train epoch to MLflow."""
+    if mlflow:
+        mlflow.log_metrics(
+            metrics={
+                **sanitize_dict(trainer.lr),
+                **sanitize_dict(trainer.label_loss_items(trainer.tloss, prefix="train")),
+            },
+            step=trainer.epoch,
+        )
+
+
+def on_fit_epoch_end(trainer):
+    """Log training metrics at the end of each fit epoch to MLflow."""
+    if mlflow:
+        mlflow.log_metrics(metrics=sanitize_dict(trainer.metrics), step=trainer.epoch)
+
+
+def on_train_end(trainer):
+    """Log model artifacts at the end of training."""
+    if not mlflow:
+        return
+    mlflow.log_artifact(str(trainer.best.parent))  # log save_dir/weights directory with best.pt and last.pt
+    for f in trainer.save_dir.glob("*"):  # log all other files in save_dir
+        if f.suffix in {".png", ".jpg", ".csv", ".pt", ".yaml"}:
+            mlflow.log_artifact(str(f))
+    keep_run_active = os.environ.get("MLFLOW_KEEP_RUN_ACTIVE", "False").lower() == "true"
+    if keep_run_active:
+        LOGGER.info(f"{PREFIX}mlflow run still alive, remember to close it using mlflow.end_run()")
+    else:
+        mlflow.end_run()
+        LOGGER.debug(f"{PREFIX}mlflow run ended")
+
+    LOGGER.info(
+        f"{PREFIX}results logged to {mlflow.get_tracking_uri()}\n{PREFIX}disable with 'yolo settings mlflow=False'"
+    )
+
+
+callbacks = (
+    {
+        "on_pretrain_routine_end": on_pretrain_routine_end,
+        "on_train_epoch_end": on_train_epoch_end,
+        "on_fit_epoch_end": on_fit_epoch_end,
+        "on_train_end": on_train_end,
+    }
+    if mlflow
+    else {}
+)

ultralytics/utils/callbacks/neptune.py

@@ -0,0 +1,126 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from ultralytics.utils import LOGGER, SETTINGS, TESTS_RUNNING
+
+try:
+    assert not TESTS_RUNNING  # do not log pytest
+    assert SETTINGS["neptune"] is True  # verify integration is enabled
+
+    import neptune
+    from neptune.types import File
+
+    assert hasattr(neptune, "__version__")
+
+    run = None  # NeptuneAI experiment logger instance
+
+except (ImportError, AssertionError):
+    neptune = None
+
+
+def _log_scalars(scalars: dict, step: int = 0) -> None:
+    """Log scalars to the NeptuneAI experiment logger.
+
+    Args:
+        scalars (dict): Dictionary of scalar values to log to NeptuneAI.
+        step (int, optional): The current step or iteration number for logging.
+
+    Examples:
+        >>> metrics = {"mAP": 0.85, "loss": 0.32}
+        >>> _log_scalars(metrics, step=100)
+    """
+    if run:
+        for k, v in scalars.items():
+            run[k].append(value=v, step=step)
+
+
+def _log_images(imgs_dict: dict, group: str = "") -> None:
+    """Log images to the NeptuneAI experiment logger.
+
+    This function logs image data to Neptune.ai when a valid Neptune run is active. Images are organized under the
+    specified group name.
+
+    Args:
+        imgs_dict (dict): Dictionary of images to log, with keys as image names and values as image data.
+        group (str, optional): Group name to organize images under in the Neptune UI.
+
+    Examples:
+        >>> # Log validation images
+        >>> _log_images({"val_batch": img_tensor}, group="validation")
+    """
+    if run:
+        for k, v in imgs_dict.items():
+            run[f"{group}/{k}"].upload(File(v))
+
+
+def _log_plot(title: str, plot_path: str) -> None:
+    """Log plots to the NeptuneAI experiment logger."""
+    import matplotlib.image as mpimg
+    import matplotlib.pyplot as plt
+
+    img = mpimg.imread(plot_path)
+    fig = plt.figure()
+    ax = fig.add_axes([0, 0, 1, 1], frameon=False, aspect="auto", xticks=[], yticks=[])  # no ticks
+    ax.imshow(img)
+    run[f"Plots/{title}"].upload(fig)
+
+
+def on_pretrain_routine_start(trainer) -> None:
+    """Initialize NeptuneAI run and log hyperparameters before training starts."""
+    try:
+        global run
+        run = neptune.init_run(
+            project=trainer.args.project or "Ultralytics",
+            name=trainer.args.name,
+            tags=["Ultralytics"],
+        )
+        run["Configuration/Hyperparameters"] = {k: "" if v is None else v for k, v in vars(trainer.args).items()}
+    except Exception as e:
+        LOGGER.warning(f"NeptuneAI installed but not initialized correctly, not logging this run. {e}")
+
+
+def on_train_epoch_end(trainer) -> None:
+    """Log training metrics and learning rate at the end of each training epoch."""
+    _log_scalars(trainer.label_loss_items(trainer.tloss, prefix="train"), trainer.epoch + 1)
+    _log_scalars(trainer.lr, trainer.epoch + 1)
+    if trainer.epoch == 1:
+        _log_images({f.stem: str(f) for f in trainer.save_dir.glob("train_batch*.jpg")}, "Mosaic")
+
+
+def on_fit_epoch_end(trainer) -> None:
+    """Log model info and validation metrics at the end of each fit epoch."""
+    if run and trainer.epoch == 0:
+        from ultralytics.utils.torch_utils import model_info_for_loggers
+
+        run["Configuration/Model"] = model_info_for_loggers(trainer)
+    _log_scalars(trainer.metrics, trainer.epoch + 1)
+
+
+def on_val_end(validator) -> None:
+    """Log validation images at the end of validation."""
+    if run:
+        # Log val_labels and val_pred
+        _log_images({f.stem: str(f) for f in validator.save_dir.glob("val*.jpg")}, "Validation")
+
+
+def on_train_end(trainer) -> None:
+    """Log final results, plots, and model weights at the end of training."""
+    if run:
+        # Log final results, CM matrix + PR plots
+        for f in [*trainer.plots.keys(), *trainer.validator.plots.keys()]:
+            if "batch" not in f.name:
+                _log_plot(title=f.stem, plot_path=f)
+        # Log the final model
+        run[f"weights/{trainer.args.name or trainer.args.task}/{trainer.best.name}"].upload(File(str(trainer.best)))
+
+
+callbacks = (
+    {
+        "on_pretrain_routine_start": on_pretrain_routine_start,
+        "on_train_epoch_end": on_train_epoch_end,
+        "on_fit_epoch_end": on_fit_epoch_end,
+        "on_val_end": on_val_end,
+        "on_train_end": on_train_end,
+    }
+    if neptune
+    else {}
+)

ultralytics/utils/callbacks/platform.py

@@ -0,0 +1,498 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import os
+import platform
+import re
+import socket
+import sys
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from time import time
+
+from ultralytics.utils import ENVIRONMENT, GIT, LOGGER, PYTHON_VERSION, RANK, SETTINGS, TESTS_RUNNING, Retry, colorstr
+
+PREFIX = colorstr("Platform: ")
+
+# Configurable platform URL for debugging (e.g. ULTRALYTICS_PLATFORM_URL=http://localhost:3000)
+PLATFORM_URL = os.getenv("ULTRALYTICS_PLATFORM_URL", "https://platform.ultralytics.com").rstrip("/")
+PLATFORM_API_URL = f"{PLATFORM_URL}/api/webhooks"
+
+
+def slugify(text):
+    """Convert text to URL-safe slug (e.g., 'My Project 1' -> 'my-project-1')."""
+    if not text:
+        return text
+    return re.sub(r"-+", "-", re.sub(r"[^a-z0-9\s-]", "", str(text).lower()).replace(" ", "-")).strip("-")[:128]
+
+
+try:
+    assert not TESTS_RUNNING  # do not log pytest
+    assert SETTINGS.get("platform", False) is True or os.getenv("ULTRALYTICS_API_KEY") or SETTINGS.get("api_key")
+    _api_key = os.getenv("ULTRALYTICS_API_KEY") or SETTINGS.get("api_key")
+    assert _api_key  # verify API key is present
+
+    import requests
+
+    from ultralytics.utils.logger import ConsoleLogger, SystemLogger
+    from ultralytics.utils.torch_utils import model_info_for_loggers
+
+    _executor = ThreadPoolExecutor(max_workers=10)  # Bounded thread pool for async operations
+
+except (AssertionError, ImportError):
+    _api_key = None
+
+
+def resolve_platform_uri(uri, hard=True):
+    """Resolve ul:// URIs to signed URLs by authenticating with Ultralytics Platform.
+
+    Formats:
+        ul://username/datasets/slug  -> Returns signed URL to NDJSON file
+        ul://username/project/model  -> Returns signed URL to .pt file
+
+    Args:
+        uri (str): Platform URI starting with "ul://".
+        hard (bool): Whether to raise an error if resolution fails.
+
+    Returns:
+        (str | None): Signed URL on success, None if not found and hard=False.
+
+    Raises:
+        ValueError: If API key is missing/invalid or URI format is wrong.
+        PermissionError: If access is denied.
+        RuntimeError: If resource is not ready (e.g., dataset still processing).
+        FileNotFoundError: If resource not found and hard=True.
+        ConnectionError: If network request fails and hard=True.
+    """
+    import requests
+
+    path = uri[5:]  # Remove "ul://"
+    parts = path.split("/")
+
+    api_key = os.getenv("ULTRALYTICS_API_KEY") or SETTINGS.get("api_key")
+    if not api_key:
+        raise ValueError(f"ULTRALYTICS_API_KEY required for '{uri}'. Get key at {PLATFORM_URL}/settings")
+
+    base = PLATFORM_API_URL
+    headers = {"Authorization": f"Bearer {api_key}"}
+
+    # ul://username/datasets/slug
+    if len(parts) == 3 and parts[1] == "datasets":
+        username, _, slug = parts
+        url = f"{base}/datasets/{username}/{slug}/export"
+
+    # ul://username/project/model
+    elif len(parts) == 3:
+        username, project, model = parts
+        url = f"{base}/models/{username}/{project}/{model}/download"
+
+    else:
+        raise ValueError(f"Invalid platform URI: {uri}. Use ul://user/datasets/name or ul://user/project/model")
+
+    try:
+        timeout = 3600 if "/datasets/" in url else 90  # NDJSON generation can be slow for large datasets
+        r = requests.head(url, headers=headers, allow_redirects=False, timeout=timeout)
+
+        # Handle redirect responses (301, 302, 303, 307, 308)
+        if 300 <= r.status_code < 400 and "location" in r.headers:
+            return r.headers["location"]  # Return signed URL
+
+        # Handle error responses
+        if r.status_code == 401:
+            raise ValueError(f"Invalid ULTRALYTICS_API_KEY for '{uri}'")
+        if r.status_code == 403:
+            raise PermissionError(f"Access denied for '{uri}'. Check dataset/model visibility settings.")
+        if r.status_code == 404:
+            if hard:
+                raise FileNotFoundError(f"Not found on platform: {uri}")
+            LOGGER.warning(f"Not found on platform: {uri}")
+            return None
+        if r.status_code == 409:
+            raise RuntimeError(f"Resource not ready: {uri}. Dataset may still be processing.")
+
+        # Unexpected response
+        r.raise_for_status()
+        raise RuntimeError(f"Unexpected response from platform for '{uri}': {r.status_code}")
+
+    except requests.exceptions.RequestException as e:
+        if hard:
+            raise ConnectionError(f"Failed to resolve {uri}: {e}") from e
+        LOGGER.warning(f"Failed to resolve {uri}: {e}")
+        return None
+
+
+def _interp_plot(plot, n=101):
+    """Interpolate plot curve data to n points to reduce storage size."""
+    import numpy as np
+
+    if not plot.get("x") or not plot.get("y"):
+        return plot  # No interpolation needed (e.g., confusion_matrix)
+
+    x, y = np.array(plot["x"]), np.array(plot["y"])
+    if len(x) <= n:
+        return plot  # Already small enough
+
+    # New x values (101 points gives clean 0.01 increments: 0, 0.01, 0.02, ..., 1.0)
+    x_new = np.linspace(x[0], x[-1], n)
+
+    # Interpolate y values (handle both 1D and 2D arrays)
+    if y.ndim == 1:
+        y_new = np.interp(x_new, x, y)
+    else:
+        y_new = np.array([np.interp(x_new, x, yi) for yi in y])
+
+    # Also interpolate ap if present (for PR curves)
+    result = {**plot, "x": x_new.tolist(), "y": y_new.tolist()}
+    if "ap" in plot:
+        result["ap"] = plot["ap"]  # Keep AP values as-is (per-class scalars)
+
+    return result
+
+
+def _send(event, data, project, name, model_id=None, retry=2):
+    """Send event to Platform endpoint with retry logic."""
+    payload = {"event": event, "project": project, "name": name, "data": data}
+    if model_id:
+        payload["modelId"] = model_id
+
+    @Retry(times=retry, delay=1)
+    def post():
+        r = requests.post(
+            f"{PLATFORM_API_URL}/training/metrics",
+            json=payload,
+            headers={"Authorization": f"Bearer {_api_key}"},
+            timeout=30,
+        )
+        if 400 <= r.status_code < 500 and r.status_code not in {408, 429}:
+            LOGGER.warning(f"{PREFIX}Failed to send {event}: {r.status_code} {r.reason}")
+            return None  # Don't retry client errors (except 408 timeout, 429 rate limit)
+        r.raise_for_status()
+        return r.json()
+
+    try:
+        return post()
+    except Exception as e:
+        LOGGER.debug(f"{PREFIX}Failed to send {event}: {e}")
+        return None
+
+
+def _send_async(event, data, project, name, model_id=None):
+    """Send event asynchronously using bounded thread pool."""
+    _executor.submit(_send, event, data, project, name, model_id)
+
+
+def _upload_model(model_path, project, name, progress=False, retry=1, model_id=None):
+    """Upload model checkpoint to Platform via signed URL."""
+    from ultralytics.utils.uploads import safe_upload
+
+    model_path = Path(model_path)
+    if not model_path.exists():
+        LOGGER.warning(f"{PREFIX}Model file not found: {model_path}")
+        return None
+
+    # Get signed upload URL from Platform (server sanitizes filename for storage safety)
+    @Retry(times=3, delay=2)
+    def get_signed_url():
+        payload = {"project": project, "name": name, "filename": model_path.name}
+        if model_id:
+            payload["modelId"] = model_id  # Direct lookup avoids slug mismatch from auto-increment
+        r = requests.post(
+            f"{PLATFORM_API_URL}/models/upload",
+            json=payload,
+            headers={"Authorization": f"Bearer {_api_key}"},
+            timeout=30,
+        )
+        r.raise_for_status()
+        return r.json()
+
+    try:
+        data = get_signed_url()
+    except Exception as e:
+        LOGGER.warning(f"{PREFIX}Failed to get upload URL: {e}")
+        return None
+
+    # Upload to GCS using safe_upload with retry logic and optional progress bar
+    if safe_upload(file=model_path, url=data["uploadUrl"], retry=retry, progress=progress):
+        return data.get("gcsPath")
+    return None
+
+
+def _upload_model_async(model_path, project, name, model_id=None):
+    """Upload model asynchronously using bounded thread pool."""
+    _executor.submit(_upload_model, model_path, project, name, model_id=model_id)
+
+
+def _get_environment_info():
+    """Collect comprehensive environment info using existing ultralytics utilities."""
+    import shutil
+
+    import psutil
+    import torch
+
+    from ultralytics import __version__
+    from ultralytics.utils.torch_utils import get_cpu_info, get_gpu_info
+
+    # Get RAM and disk totals
+    memory = psutil.virtual_memory()
+    disk_usage = shutil.disk_usage("/")
+
+    env = {
+        "ultralyticsVersion": __version__,
+        "hostname": socket.gethostname(),
+        "os": platform.platform(),
+        "environment": ENVIRONMENT,
+        "pythonVersion": PYTHON_VERSION,
+        "pythonExecutable": sys.executable,
+        "cpuCount": os.cpu_count() or 0,
+        "cpu": get_cpu_info(),
+        "command": " ".join(sys.argv),
+        "totalRamGb": round(memory.total / (1 << 30), 1),  # Total RAM in GB
+        "totalDiskGb": round(disk_usage.total / (1 << 30), 1),  # Total disk in GB
+    }
+
+    # Git info using cached GIT singleton (no subprocess calls)
+    try:
+        if GIT.is_repo:
+            if GIT.origin:
+                env["gitRepository"] = GIT.origin
+            if GIT.branch:
+                env["gitBranch"] = GIT.branch
+            if GIT.commit:
+                env["gitCommit"] = GIT.commit[:12]  # Short hash
+    except Exception:
+        pass
+
+    # GPU info
+    try:
+        if torch.cuda.is_available():
+            env["gpuCount"] = torch.cuda.device_count()
+            env["gpuType"] = get_gpu_info(0) if torch.cuda.device_count() > 0 else None
+    except Exception:
+        pass
+
+    return env
+
+
+def _get_project_name(trainer):
+    """Get slugified project and name from trainer args."""
+    raw = str(trainer.args.project)
+    parts = raw.split("/", 1)
+    project = f"{parts[0]}/{slugify(parts[1])}" if len(parts) == 2 else slugify(raw)
+    return project, slugify(str(trainer.args.name or "train"))
+
+
+def on_pretrain_routine_start(trainer):
+    """Initialize Platform logging at training start."""
+    if RANK not in {-1, 0} or not trainer.args.project:
+        return
+
+    project, name = _get_project_name(trainer)
+    LOGGER.info(f"{PREFIX}Streaming training metrics to Platform")
+
+    # Single dict for all platform callback state (like trainer.hub_session for HUB callbacks)
+    ctx = {"model_id": None, "last_upload": time(), "cancelled": False, "console_logger": None, "system_logger": None}
+    trainer.platform = ctx
+
+    # Create callback to send console output to Platform
+    def send_console_output(content, line_count, chunk_id):
+        """Send batched console output to Platform webhook."""
+        _send_async(
+            "console_output",
+            {"chunkId": chunk_id, "content": content, "lineCount": line_count},
+            project,
+            name,
+            ctx["model_id"],
+        )
+
+    # Start console capture with batching (5 lines or 5 seconds)
+    ctx["console_logger"] = ConsoleLogger(batch_size=5, flush_interval=5.0, on_flush=send_console_output)
+    ctx["console_logger"].start_capture()
+
+    # Collect environment info (W&B-style metadata)
+    environment = _get_environment_info()
+
+    # Build trainArgs - callback runs before get_dataset() so args.data is still original (e.g., ul:// URIs)
+    # Note: model_info is sent later in on_fit_epoch_end (epoch 0) when the model is actually loaded
+    train_args = {k: str(v) for k, v in vars(trainer.args).items()}
+
+    # Send synchronously to get modelId for subsequent webhooks (critical, more retries)
+    response = _send(
+        "training_started",
+        {
+            "trainArgs": train_args,
+            "epochs": trainer.epochs,
+            "device": str(trainer.device),
+            "environment": environment,
+        },
+        project,
+        name,
+        retry=4,
+    )
+    if response and response.get("modelId"):
+        ctx["model_id"] = response["modelId"]
+        # Server returns actual slug (may differ from requested name due to auto-increment, e.g. "train" → "train-2")
+        if response.get("modelSlug"):
+            ctx["model_slug"] = response["modelSlug"]
+            url = f"{PLATFORM_URL}/{project}/{ctx['model_slug']}"
+            LOGGER.info(f"{PREFIX}View model at {url}")
+        # Check for immediate cancellation (cancelled before training started)
+        # Note: trainer.stop is set in on_pretrain_routine_end (after _setup_train resets it)
+        if response.get("cancelled"):
+            ctx["cancelled"] = True
+    else:
+        LOGGER.warning(f"{PREFIX}Failed to register training session - metrics may not sync to Platform")
+
+
+def on_pretrain_routine_end(trainer):
+    """Apply pre-start cancellation after _setup_train resets trainer.stop."""
+    ctx = getattr(trainer, "platform", None)
+    if ctx and ctx["cancelled"]:
+        LOGGER.info(f"{PREFIX}Training cancelled from Platform before starting ✅")
+        trainer.stop = True
+
+
+def on_fit_epoch_end(trainer):
+    """Log training and system metrics at epoch end."""
+    ctx = getattr(trainer, "platform", None)
+    if not ctx or RANK not in {-1, 0} or not trainer.args.project:
+        return
+
+    project, name = _get_project_name(trainer)
+    metrics = {**trainer.label_loss_items(trainer.tloss, prefix="train"), **trainer.metrics}
+
+    if trainer.optimizer and trainer.optimizer.param_groups:
+        metrics["lr"] = trainer.optimizer.param_groups[0]["lr"]
+
+    # Extract model info at epoch 0 (sent as separate field, not in metrics)
+    model_info = None
+    if trainer.epoch == 0:
+        try:
+            info = model_info_for_loggers(trainer)
+            model_info = {
+                "parameters": info.get("model/parameters", 0),
+                "gflops": info.get("model/GFLOPs", 0),
+                "speedMs": info.get("model/speed_PyTorch(ms)", 0),
+            }
+        except Exception:
+            pass
+
+    # Get system metrics (cache SystemLogger in platform context for efficiency)
+    system = {}
+    try:
+        if not ctx["system_logger"]:
+            ctx["system_logger"] = SystemLogger()
+        system = ctx["system_logger"].get_metrics(rates=True)
+    except Exception:
+        pass
+
+    payload = {
+        "epoch": trainer.epoch,
+        "metrics": metrics,
+        "system": system,
+        "fitness": trainer.fitness,
+        "best_fitness": trainer.best_fitness,
+    }
+    if model_info:
+        payload["modelInfo"] = model_info
+
+    def _send_and_check_cancel():
+        """Send epoch_end and check response for cancellation (runs in background thread)."""
+        response = _send("epoch_end", payload, project, name, ctx["model_id"], retry=1)
+        if response and response.get("cancelled"):
+            LOGGER.info(f"{PREFIX}Training cancelled from Platform ✅")
+            trainer.stop = True
+            ctx["cancelled"] = True
+
+    _executor.submit(_send_and_check_cancel)
+
+
+def on_model_save(trainer):
+    """Upload model checkpoint (rate limited to every 15 min)."""
+    ctx = getattr(trainer, "platform", None)
+    if not ctx or RANK not in {-1, 0} or not trainer.args.project:
+        return
+
+    # Rate limit to every 15 minutes (900 seconds)
+    if time() - ctx["last_upload"] < 900:
+        return
+
+    model_path = trainer.best if trainer.best and Path(trainer.best).exists() else trainer.last
+    if not model_path:
+        return
+
+    project, name = _get_project_name(trainer)
+    _upload_model_async(model_path, project, name, model_id=ctx["model_id"])
+    ctx["last_upload"] = time()
+
+
+def on_train_end(trainer):
+    """Log final results, upload best model, and send validation plot data."""
+    ctx = getattr(trainer, "platform", None)
+    if not ctx or RANK not in {-1, 0} or not trainer.args.project:
+        return
+
+    project, name = _get_project_name(trainer)
+
+    if ctx["cancelled"]:
+        LOGGER.info(f"{PREFIX}Uploading partial results for cancelled training")
+
+    # Stop console capture
+    if ctx["console_logger"]:
+        ctx["console_logger"].stop_capture()
+        ctx["console_logger"] = None
+
+    # Upload best model (blocking with progress bar to ensure it completes)
+    gcs_path = None
+    model_size = None
+    if trainer.best and Path(trainer.best).exists():
+        model_size = Path(trainer.best).stat().st_size
+        gcs_path = _upload_model(trainer.best, project, name, progress=True, retry=3, model_id=ctx["model_id"])
+        if not gcs_path:
+            LOGGER.warning(f"{PREFIX}Model will not be available for download on Platform (upload failed)")
+
+    # Collect plots from trainer and validator, deduplicating by type
+    plots_by_type = {}
+    for info in getattr(trainer, "plots", {}).values():
+        if info.get("data") and info["data"].get("type"):
+            plots_by_type[info["data"]["type"]] = info["data"]
+    for info in getattr(getattr(trainer, "validator", None), "plots", {}).values():
+        if info.get("data") and info["data"].get("type"):
+            plots_by_type.setdefault(info["data"]["type"], info["data"])  # Don't overwrite trainer plots
+    plots = [_interp_plot(p) for p in plots_by_type.values()]  # Interpolate curves to reduce size
+
+    # Get class names
+    names = getattr(getattr(trainer, "validator", None), "names", None) or (trainer.data or {}).get("names")
+    class_names = list(names.values()) if isinstance(names, dict) else list(names) if names else None
+
+    _send(
+        "training_complete",
+        {
+            "results": {
+                "metrics": {**trainer.metrics, "fitness": trainer.fitness},
+                "bestEpoch": getattr(trainer, "best_epoch", trainer.epoch),
+                "bestFitness": trainer.best_fitness,
+                "modelPath": gcs_path,  # Only send GCS path, not local path
+                "modelSize": model_size,
+            },
+            "classNames": class_names,
+            "plots": plots,
+        },
+        project,
+        name,
+        ctx["model_id"],
+        retry=4,  # Critical, more retries
+    )
+    url = f"{PLATFORM_URL}/{project}/{ctx.get('model_slug', name)}"
+    LOGGER.info(f"{PREFIX}View results at {url}")
+
+
+callbacks = (
+    {
+        "on_pretrain_routine_start": on_pretrain_routine_start,
+        "on_pretrain_routine_end": on_pretrain_routine_end,
+        "on_fit_epoch_end": on_fit_epoch_end,
+        "on_model_save": on_model_save,
+        "on_train_end": on_train_end,
+    }
+    if _api_key
+    else {}
+)

ultralytics/utils/callbacks/raytune.py

@@ -0,0 +1,42 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from ultralytics.utils import SETTINGS
+
+try:
+    assert SETTINGS["raytune"] is True  # verify integration is enabled
+    import ray
+    from ray import tune
+    from ray.air import session
+
+except (ImportError, AssertionError):
+    tune = None
+
+
+def on_fit_epoch_end(trainer):
+    """Report training metrics to Ray Tune at epoch end when a Ray session is active.
+
+    Captures metrics from the trainer object and sends them to Ray Tune with the current epoch number, enabling
+    hyperparameter tuning optimization. Only executes when within an active Ray Tune session.
+
+    Args:
+        trainer (ultralytics.engine.trainer.BaseTrainer): The Ultralytics trainer object containing metrics and epochs.
+
+    Examples:
+        >>> # Called automatically by the Ultralytics training loop
+        >>> on_fit_epoch_end(trainer)
+
+    References:
+        Ray Tune docs: https://docs.ray.io/en/latest/tune/index.html
+    """
+    if ray.train._internal.session.get_session():  # check if Ray Tune session is active
+        metrics = trainer.metrics
+        session.report({**metrics, **{"epoch": trainer.epoch + 1}})
+
+
+callbacks = (
+    {
+        "on_fit_epoch_end": on_fit_epoch_end,
+    }
+    if tune
+    else {}
+)

ultralytics/utils/callbacks/tensorboard.py

@@ -0,0 +1,689 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import warnings
+
+from ultralytics.utils import DEFAULT_CFG, LOGGER, RANK, SETTINGS, TESTS_RUNNING, colorstr, torch_utils
+from ultralytics.utils.torch_utils import smart_inference_mode
+
+try:
+    assert not TESTS_RUNNING  # do not log pytest
+    assert SETTINGS["tensorboard"] is True  # verify integration is enabled
+    WRITER = None  # TensorBoard SummaryWriter instance
+    PREFIX = colorstr("TensorBoard: ")
+
+    # Imports below only required if TensorBoard enabled
+    from copy import deepcopy
+    from pathlib import Path
+
+    import cv2
+    import numpy as np
+    import torch
+    from torch.utils.tensorboard import SummaryWriter
+
+except (ImportError, AssertionError, TypeError, AttributeError):
+    # TypeError for handling 'Descriptors cannot not be created directly.' protobuf errors in Windows
+    # AttributeError: module 'tensorflow' has no attribute 'io' if 'tensorflow' not installed
+    SummaryWriter = None
+
+# Global state for validation logging
+_val_sample_indices_to_log = set()
+_val_seen_samples = 0
+_val_index_offset = 0
+_current_epoch = 0
+
+
+def _reset_val_logging_state() -> None:
+    """Reset validation image logging state."""
+    global _val_sample_indices_to_log, _val_seen_samples, _val_index_offset
+    _val_sample_indices_to_log = set()
+    _val_seen_samples = 0
+    _val_index_offset = 0
+
+
+def _close_writer() -> None:
+    """Flush and close the active TensorBoard writer if one exists."""
+    global WRITER
+    if WRITER:
+        try:
+            WRITER.flush()
+        finally:
+            WRITER.close()
+            WRITER = None
+
+
+def _sanitize_tb_tag_component(value: str, max_len: int = 80) -> str:
+    """Return a TensorBoard-friendly tag component."""
+    value = "".join(c if c.isalnum() or c in ("-", "_", ".") else "_" for c in str(value).strip())
+    value = value.strip("._")
+    return value[:max_len] if value else "unknown"
+
+
+def _create_detection_image(images, targets, paths=None, names=None, is_labels=True, conf_thres=0.25):
+    """Create detection visualization image for TensorBoard logging.
+
+    Args:
+        images (torch.Tensor): Batch of images (B, C, H, W)
+        targets (torch.Tensor | np.ndarray): Target labels or predictions
+            - For labels: (N, 6) [batch_idx, class_id, x, y, w, h] (normalized xywh)
+            - For predictions: (N, 7) [batch_idx, class_id, x, y, w, h, conf] (normalized xywh)
+        paths (list): Image paths
+        names (dict | list): Class names
+        is_labels (bool): Whether targets are ground truth labels (True) or predictions (False)
+        conf_thres (float): Confidence threshold for filtering predictions
+
+    Returns:
+        np.ndarray: RGB image with bounding boxes drawn
+    """
+    if not WRITER:
+        return None
+
+    try:
+        from ultralytics.utils.plotting import Annotator, colors
+        from ultralytics.utils.ops import xywh2xyxy
+        import math
+        from pathlib import Path
+
+        # Convert to numpy
+        if isinstance(images, torch.Tensor):
+            images = images.cpu().float().numpy()
+        if isinstance(targets, torch.Tensor):
+            targets = targets.cpu().numpy()
+
+        max_size = 1920  # max image size
+        max_subplots = 1  # limit to 1 image for cleaner visualization
+        bs, _, h, w = images.shape  # batch size, _, height, width
+        bs = min(bs, max_subplots)
+        ns = np.ceil(bs**0.5)  # number of subplots (square)
+
+        if np.max(images[0]) <= 1:
+            images = images * 255  # de-normalise
+
+        # Build Image
+        mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8)
+        for i, im in enumerate(images):
+            if i == max_subplots:
+                break
+            x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
+            im = im.transpose(1, 2, 0)  # CHW to HWC
+            # Images are in RGB format (converted from BGR in augmentation pipeline with default bgr=0.0)
+            # No color conversion needed for TensorBoard
+            mosaic[y : y + h, x : x + w, :] = im.astype(np.uint8)
+
+        # Resize (optional)
+        scale = max_size / ns / max(h, w)
+        if scale < 1:
+            h = math.ceil(scale * h)
+            w = math.ceil(scale * w)
+            mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)), interpolation=cv2.INTER_AREA)
+
+        # Annotate with thin lines
+        fs = int((h + w) * ns * 0.01)  # font size
+        line_width = max(1, round(fs / 30))  # thin line width
+        annotator = Annotator(mosaic, line_width=line_width, font_size=fs, pil=True)
+
+        for i in range(bs):
+            x, y = int(w * (i // ns)), int(h * (i % ns))  # block origin
+            annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2)  # borders
+            if paths:
+                annotator.text([x + 5, y + 5], text=Path(paths[i]).name[:40], txt_color=(220, 220, 220))
+
+            if len(targets) > 0:
+                ti = targets[targets[:, 0] == i]  # image targets
+                boxes = xywh2xyxy(ti[:, 2:6]).T
+                classes = ti[:, 1].astype("int")
+                conf = None if is_labels else ti[:, 6]  # confidence only for predictions
+
+                if boxes.shape[1]:
+                    if boxes.max() <= 1.01:  # if normalized with tolerance 0.01
+                        boxes[[0, 2]] *= w  # scale to pixels
+                        boxes[[1, 3]] *= h
+                    elif scale < 1:  # absolute coords need scale if image scales
+                        boxes *= scale
+                boxes[[0, 2]] += x
+                boxes[[1, 3]] += y
+
+                for j, box in enumerate(boxes.T.tolist()):
+                    cls = classes[j]
+                    color = colors(cls)
+
+                    # Get class name
+                    if names:
+                        if isinstance(names, dict):
+                            cls_name = names.get(cls, str(cls))
+                        elif isinstance(names, list):
+                            cls_name = names[cls] if cls < len(names) else str(cls)
+                        else:
+                            cls_name = str(cls)
+                    else:
+                        cls_name = str(cls)
+
+                    # For predictions, show confidence score; for labels, show class name
+                    if is_labels or (conf is not None and conf[j] > conf_thres):
+                        if not is_labels and conf is not None:
+                            label = f"{cls_name} {conf[j]:.2f}"  # Show class and confidence for predictions
+                        else:
+                            label = cls_name  # Show class name for ground truth
+                        annotator.box_label(box, label, color=color)
+
+        # Get result from Annotator (returns RGB when pil=True)
+        result = np.asarray(annotator.result())
+        return result
+
+    except Exception as e:
+        LOGGER.warning(f"{PREFIX}Failed to create detection image: {e}")
+        return None
+
+
+def _output_to_target(output, max_det=300):
+    """Convert model output to target format for plotting.
+
+    Args:
+        output (list): List of prediction dicts or tensors
+            - For dict format: [{'bboxes': (N,4), 'conf': (N,), 'cls': (N,), ...}, ...]
+            - For tensor format: [(N, 6) [x1, y1, x2, y2, conf, cls], ...]
+        max_det (int): Maximum detections per image
+
+    Returns:
+        np.ndarray: (M, 7) [batch_idx, class_id, x, y, w, h, conf] (normalized xywh)
+    """
+    try:
+        from ultralytics.utils.ops import xyxy2xywh
+
+        targets = []
+        for i, o in enumerate(output):
+            # Handle dict format (new postprocess output)
+            if isinstance(o, dict):
+                bboxes = o.get('bboxes', torch.empty((0, 4)))
+                conf = o.get('conf', torch.empty((0,)))
+                cls = o.get('cls', torch.empty((0,)))
+
+                if len(bboxes) == 0:
+                    continue
+
+                # Limit detections
+                if len(bboxes) > max_det:
+                    bboxes = bboxes[:max_det]
+                    conf = conf[:max_det]
+                    cls = cls[:max_det]
+
+                # Ensure proper shapes
+                if conf.dim() == 1:
+                    conf = conf.unsqueeze(-1)
+                if cls.dim() == 1:
+                    cls = cls.unsqueeze(-1)
+
+                # Create batch index
+                j = torch.full((len(bboxes), 1), i, device=bboxes.device)
+
+                # Concatenate: [batch_idx, cls, xywh, conf]
+                targets.append(torch.cat((j, cls, xyxy2xywh(bboxes), conf), 1))
+
+            # Handle tensor format (legacy)
+            else:
+                if len(o) == 0:
+                    continue
+                o = o[:max_det]  # limit detections
+                box, conf, cls = o[:, :4], o[:, 4:5], o[:, 5:6]
+                j = torch.full((conf.shape[0], 1), i, device=o.device)  # batch index
+                targets.append(torch.cat((j, cls, xyxy2xywh(box), conf), 1))
+
+        if targets:
+            return torch.cat(targets, 0).cpu().numpy()
+        else:
+            return np.empty((0, 7))
+
+    except Exception as e:
+        LOGGER.warning(f"{PREFIX}Failed to convert output to target: {e}")
+        return np.empty((0, 7))
+
+
+def _log_scalars(scalars: dict | None, step: int = 0) -> None:
+    """Log scalar values to TensorBoard.
+
+    Args:
+        scalars (dict | None): Dictionary of scalar values to log to TensorBoard. Keys are scalar names and values are
+            the corresponding scalar values. If None or empty, nothing is logged.
+        step (int): Global step value to record with the scalar values. Used for x-axis in TensorBoard graphs.
+
+    Examples:
+        Log training metrics
+        >>> metrics = {"loss": 0.5, "accuracy": 0.95}
+        >>> _log_scalars(metrics, step=100)
+    """
+    if WRITER and scalars:
+        for k, v in scalars.items():
+            WRITER.add_scalar(k, v, step)
+
+
+@smart_inference_mode()
+def _log_tensorboard_graph(trainer) -> None:
+    """Log model graph to TensorBoard.
+
+    This function attempts to visualize the model architecture in TensorBoard by tracing the model with a dummy input
+    tensor. It first tries a simple method suitable for YOLO models, and if that fails, falls back to a more complex
+    approach for models like RTDETR that may require special handling.
+
+    Args:
+        trainer (ultralytics.engine.trainer.BaseTrainer): The trainer object containing the model to visualize. Must
+            have attributes model and args with imgsz.
+
+    Notes:
+        This function requires TensorBoard integration to be enabled and the global WRITER to be initialized.
+        It handles potential warnings from the PyTorch JIT tracer and attempts to gracefully handle different
+        model architectures.
+    """
+    # Input image - use CPU to avoid corrupting CUDA context if tracing fails
+    imgsz = trainer.args.imgsz
+    imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz
+    im = torch.zeros((1, 3, *imgsz), device="cpu", dtype=torch.float32)
+
+    # Try simple method first (YOLO)
+    try:
+        model = deepcopy(torch_utils.unwrap_model(trainer.model)).cpu().eval()
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore",
+                message="The input to trace is already a ScriptModule, tracing it is a no-op.*",
+                category=UserWarning,
+            )
+            WRITER.add_graph(torch.jit.trace(model, im, strict=False), [])
+        LOGGER.info(f"{PREFIX}model graph visualization added ✅")
+        return
+    except Exception as e1:
+        # Fallback to TorchScript export steps (RTDETR)
+        try:
+            model = deepcopy(torch_utils.unwrap_model(trainer.model)).cpu().eval()
+            model = model.fuse(verbose=False)
+            for m in model.modules():
+                if hasattr(m, "export"):  # Detect, RTDETRDecoder (Segment and Pose use Detect base class)
+                    m.export = True
+                    m.format = "torchscript"
+            model(im)  # dry run
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    "ignore",
+                    message="The input to trace is already a ScriptModule, tracing it is a no-op.*",
+                    category=UserWarning,
+                )
+                WRITER.add_graph(torch.jit.trace(model, im, strict=False), [])
+            LOGGER.info(f"{PREFIX}model graph visualization added ✅")
+        except Exception as e2:
+            LOGGER.warning(f"{PREFIX}TensorBoard graph visualization failure: {e1} -> {e2}")
+
+
+def on_pretrain_routine_start(trainer) -> None:
+    """Initialize TensorBoard logging with SummaryWriter."""
+    if SummaryWriter:
+        try:
+            global WRITER
+            WRITER = SummaryWriter(str(trainer.save_dir))
+            LOGGER.info(f"{PREFIX}Start with 'tensorboard --logdir {trainer.save_dir}', view at http://localhost:6006/")
+        except Exception as e:
+            LOGGER.warning(f"{PREFIX}TensorBoard not initialized correctly, not logging this run. {e}")
+
+
+def on_train_start(trainer) -> None:
+    """Log TensorBoard graph."""
+    if WRITER:
+        _log_tensorboard_graph(trainer)
+
+
+def on_train_epoch_end(trainer) -> None:
+    """Log scalar statistics at the end of a training epoch."""
+    _log_scalars(trainer.label_loss_items(trainer.tloss, prefix="train"), trainer.epoch + 1)
+    _log_scalars(trainer.lr, trainer.epoch + 1)
+
+
+def on_fit_epoch_end(trainer) -> None:
+    """Log epoch metrics at end of training epoch."""
+    _log_scalars(trainer.metrics, trainer.epoch + 1)
+
+
+def on_val_start(validator) -> None:
+    """Initialize validation logging - select samples to log."""
+    global _val_sample_indices_to_log, _val_seen_samples, _val_index_offset
+    if WRITER:
+        dataset = getattr(validator.dataloader, "dataset", None)
+        total_samples = len(dataset or [])
+        if total_samples <= 0:
+            total_samples = len(validator.dataloader)
+        num_samples = min(50, total_samples)
+        sample_indices = np.linspace(0, max(total_samples - 1, 0), num_samples, dtype=int).tolist()
+
+        sampler = getattr(validator.dataloader, "sampler", None)
+        if sampler and hasattr(sampler, "_get_rank_indices"):
+            start_idx, end_idx = sampler._get_rank_indices()
+            _val_index_offset = start_idx
+            _val_sample_indices_to_log = {idx - start_idx for idx in sample_indices if start_idx <= idx < end_idx}
+        else:
+            _val_index_offset = 0
+            _val_sample_indices_to_log = set(sample_indices)
+
+        _val_seen_samples = 0
+
+
+def on_val_batch_end(validator) -> None:
+    """Log validation sample images to TensorBoard as 2x2 grid if 3D data available."""
+    global _current_epoch, _val_sample_indices_to_log, _val_seen_samples, _val_index_offset
+
+    if not WRITER or not _val_sample_indices_to_log:
+        return
+
+    try:
+        batch = validator.batch
+        batch_size = len(batch.get("im_file", [])) or int(batch["img"].shape[0])
+        batch_start = _val_seen_samples
+        batch_end = batch_start + batch_size
+        sample_indices = [i for i in range(batch_start, batch_end) if i in _val_sample_indices_to_log]
+        _val_seen_samples = batch_end
+        if not sample_indices:
+            return
+
+        # 2x2 grid [Target 2D | Target 3D] / [Pred 2D | Pred 3D]
+        for sample_idx in sample_indices:
+            batch_sample_idx = sample_idx - batch_start
+            _log_3d_visualization(validator, batch, batch_sample_idx, sample_idx + _val_index_offset)
+
+    except Exception as e:
+        LOGGER.warning(f"{PREFIX}Failed to log validation images: {e}")
+
+
+def _draw_dashed_rectangle(img, p1, p2, color, thickness=1, dash=8, gap=5):
+    """Draw a dashed rectangle on an image."""
+    x1, y1 = p1
+    x2, y2 = p2
+    for x in range(x1, x2, dash + gap):
+        cv2.line(img, (x, y1), (min(x + dash, x2), y1), color, thickness, cv2.LINE_AA)
+        cv2.line(img, (x, y2), (min(x + dash, x2), y2), color, thickness, cv2.LINE_AA)
+    for y in range(y1, y2, dash + gap):
+        cv2.line(img, (x1, y), (x1, min(y + dash, y2)), color, thickness, cv2.LINE_AA)
+        cv2.line(img, (x2, y), (x2, min(y + dash, y2)), color, thickness, cv2.LINE_AA)
+
+
+def _draw_2d_boxes_on_image(img, boxes_xyxy, cls_ids, names=None, confs=None, thickness=1, diff_cls=None):
+    """Draw 2D bounding boxes on an image using cv2 (BGR color conventions).
+
+    Args:
+        img: (H, W, 3) uint8 image (BGR).
+        boxes_xyxy: (N, 4) array of [x1, y1, x2, y2] in pixel coords.
+        cls_ids: (N,) array of class IDs.
+        names: Dict or list of class names.
+        confs: (N,) array of confidences, or None for GT.
+        thickness: Line thickness.
+    """
+    from ultralytics.utils.plotting import colors as get_color
+
+    for i in range(len(boxes_xyxy)):
+        x1, y1, x2, y2 = [int(v) for v in boxes_xyxy[i]]
+        cls = int(cls_ids[i])
+        color = get_color(cls, bgr=True)
+        is_diff1 = diff_cls is not None and int(diff_cls[i]) == 1
+        if is_diff1:
+            _draw_dashed_rectangle(img, (x1, y1), (x2, y2), color, thickness=max(thickness, 1))
+        else:
+            cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness, cv2.LINE_AA)
+
+        # Label
+        if names:
+            cls_name = names.get(cls, str(cls)) if isinstance(names, dict) else (
+                names[cls] if isinstance(names, list) and cls < len(names) else str(cls))
+        else:
+            cls_name = str(cls)
+        label = f"{cls_name} {confs[i]:.2f}" if confs is not None else cls_name
+        if diff_cls is not None:
+            label = f"{label} d{int(diff_cls[i])}"
+        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+        cv2.rectangle(img, (x1, y1 - th - 4), (x1 + tw, y1), color, -1)
+        cv2.putText(img, label, (x1, y1 - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
+
+
+def _log_3d_visualization(validator, batch, batch_sample_idx=0, sample_idx=None):
+    """Log 2x2 grid visualization: [Target 2D | Target 3D] / [Pred 2D | Pred 3D].
+
+    GT depth is restored in decode_3d_target() using calib["depth_scale"].
+    Prediction depth is restored once upstream in Ground3DDetectionValidator.postprocess().
+    Both are projected with the original batch calibration.
+    """
+    labels_3d = batch.get("labels_3d")
+    if labels_3d is None or len(labels_3d) == 0:
+        return
+
+    try:
+        from ultralytics.utils.ops import xywh2xyxy
+        from ultralytics.utils.plotting_3d import (
+            collect_precomputed_edge_points_2d,
+            decode_3d_prediction,
+            decode_3d_target,
+            draw_3d_box,
+        )
+
+        # Get 3D class config from model
+        model = validator.trainer.model if hasattr(validator, "trainer") and validator.trainer else None
+        if model is None:
+            return
+        if hasattr(model, "module"):
+            model = model.module
+        face_3d_classes = getattr(model, "face_3d_classes", set())
+        complete_3d_classes = getattr(model, "complete_3d_classes", set())
+        if not face_3d_classes and not complete_3d_classes:
+            return
+
+        images = batch["img"]
+        batch_idx_t = batch.get("batch_idx")
+        batch_cls_t = batch.get("cls")
+        batch_bboxes = batch.get("bboxes")
+        if batch_idx_t is None or batch_cls_t is None:
+            return
+
+        batch_idx_np = batch_idx_t.cpu().numpy().ravel()
+        cls_np = batch_cls_t.cpu().numpy().ravel()
+        labels_3d_np = labels_3d.cpu().numpy()
+        edge_faces_points_2d = batch.get("edge_faces_points_2d")
+        edge_faces_valid = batch.get("edge_faces_valid")
+        edge_faces_points_2d_np = edge_faces_points_2d.cpu().numpy() if edge_faces_points_2d is not None else None
+        edge_faces_valid_np = edge_faces_valid.cpu().numpy() if edge_faces_valid is not None else None
+        face_visibility_score_thresh = float(
+            getattr(validator.args, "face_visibility_score_thresh", DEFAULT_CFG.face_visibility_score_thresh)
+        )
+        _, _, img_h, img_w = images.shape
+
+        batch_calib = batch.get("calib")
+        calib = batch_calib[batch_sample_idx] if batch_calib is not None and batch_sample_idx < len(batch_calib) else None
+        if calib is None:
+            LOGGER.warning(f"{PREFIX}No calib in batch, skipping 3D visualization")
+            return
+
+        batch_camera_mode = batch.get("camera_mode")
+        camera_mode = (
+            str(batch_camera_mode[batch_sample_idx])
+            if isinstance(batch_camera_mode, (list, tuple)) and batch_sample_idx < len(batch_camera_mode)
+            else "unknown"
+        )
+
+        epoch = validator.trainer.epoch if validator.training and validator.trainer else 0
+        global_sample_idx = sample_idx if sample_idx is not None else batch_sample_idx
+        rank_tag = f"rank{RANK if RANK >= 0 else 0}"
+        batch_im_files = batch.get("im_file")
+        image_name = None
+        if isinstance(batch_im_files, (list, tuple)) and batch_sample_idx < len(batch_im_files):
+            image_name = Path(str(batch_im_files[batch_sample_idx])).name
+        image_tag = _sanitize_tb_tag_component(Path(image_name).stem if image_name else "unknown")
+
+        # --- Prepare base image ---
+        im0 = images[batch_sample_idx].cpu().numpy().transpose(1, 2, 0)
+        im0 = np.ascontiguousarray(im0 * 255, dtype=np.uint8)
+        im0 = cv2.cvtColor(im0, cv2.COLOR_RGB2BGR)
+
+        # Create 4 copies for the 2x2 grid
+        im_gt_2d = im0.copy()
+        im_gt_3d = im0.copy()
+        im_pred_2d = im0.copy()
+        im_pred_3d = im0.copy()
+
+        line_thick = 1
+
+        # --- GT: Decode and draw boxes for selected image ---
+        mask_i = batch_idx_np == batch_sample_idx
+        if batch_bboxes is not None and mask_i.any():
+            gt_bboxes_xywh = batch_bboxes[mask_i].cpu().numpy()
+            gt_cls = cls_np[mask_i]
+            gt_bboxes_xyxy = xywh2xyxy(torch.from_numpy(gt_bboxes_xywh)).numpy()
+            gt_bboxes_xyxy[:, [0, 2]] *= img_w
+            gt_bboxes_xyxy[:, [1, 3]] *= img_h
+            _draw_2d_boxes_on_image(im_gt_2d, gt_bboxes_xyxy, gt_cls, names=validator.names, thickness=line_thick)
+
+        gt_indices = np.where(mask_i)[0]
+        for local_idx, idx in enumerate(gt_indices):
+            bbox_xyxy = gt_bboxes_xyxy[local_idx] if batch_bboxes is not None and mask_i.any() else None
+            d = decode_3d_target(
+                labels_3d_np[idx],
+                int(cls_np[idx]),
+                calib,
+                img_w,
+                img_h,
+                face_3d_classes,
+                complete_3d_classes,
+                score_thr=face_visibility_score_thresh,
+                bbox_xyxy=bbox_xyxy,
+            )
+            if d is not None and d.get("corners_3d") is not None:
+                if edge_faces_points_2d_np is not None and edge_faces_valid_np is not None and idx < len(edge_faces_points_2d_np):
+                    edge_points_2d = collect_precomputed_edge_points_2d(
+                        edge_faces_points_2d_np[idx],
+                        edge_faces_valid_np[idx],
+                        visible_face_types=d.get("visible_face_types", ()),
+                    )
+                    if edge_points_2d is not None:
+                        d = {**d, "edge_points_2d": edge_points_2d}
+                draw_3d_box(
+                    im_gt_3d,
+                    d["corners_3d"],
+                    calib,
+                    d.get("face_center_2d"),
+                    d.get("face_color"),
+                    edge_points_2d=d.get("edge_points_2d"),
+                    edge_color=(0, 255, 0),
+                    thickness=line_thick,
+                )
+
+        # --- Predictions: Decode and draw boxes for selected image ---
+        preds_3d_sel = getattr(validator, "_preds_3d_selected", None)
+        preds_edge_sel = getattr(validator, "_preds_edge_selected", None)
+        preds_diff_sel = getattr(validator, "_preds_diff_selected", None)
+        anchors_sel = getattr(validator, "_anchors_selected", None)
+        strides_sel = getattr(validator, "_strides_selected", None)
+        preds_2d = validator.pred
+
+        if preds_2d is not None and len(preds_2d) > batch_sample_idx:
+            pred_i = preds_2d[batch_sample_idx]
+            pred_bboxes = pred_i["bboxes"].cpu().numpy()
+            pred_cls_np = pred_i["cls"].cpu().numpy()
+            pred_conf_np = pred_i["conf"].cpu().numpy()
+            display_conf = max(float(getattr(validator.args, "visualize_conf", 0.25)), float(getattr(validator.args, "conf", 0.0)))
+            display_mask = pred_conf_np >= display_conf
+
+            if display_mask.any():
+                pred_diff_np = None
+                if preds_diff_sel is not None:
+                    diff_logits = preds_diff_sel[batch_sample_idx]
+                    if diff_logits.shape[-1] == 1:
+                        diff_prob = diff_logits[..., 0].sigmoid()
+                        diff_thres = float(getattr(validator.args, "diff_thres", 0.7))
+                        pred_diff_np = (diff_prob >= diff_thres).cpu().numpy().astype(np.int64)
+                    elif diff_logits.shape[-1] == 2:
+                        pred_diff_np = diff_logits.softmax(-1).argmax(-1).cpu().numpy().astype(np.int64)
+                _draw_2d_boxes_on_image(
+                    im_pred_2d,
+                    pred_bboxes[display_mask],
+                    pred_cls_np[display_mask],
+                    names=validator.names,
+                    confs=pred_conf_np[display_mask],
+                    thickness=line_thick,
+                    diff_cls=pred_diff_np[display_mask] if pred_diff_np is not None else None,
+                )
+
+                if preds_3d_sel is not None and anchors_sel is not None and strides_sel is not None:
+                    p3d = preds_3d_sel[batch_sample_idx].cpu().numpy()
+                    pedge = preds_edge_sel[batch_sample_idx].cpu().numpy() if preds_edge_sel is not None else None
+                    anchors_np = anchors_sel[batch_sample_idx].cpu().numpy()
+                    strides_np = strides_sel[batch_sample_idx].cpu().numpy()
+                    for i in np.where(display_mask)[0]:
+                        d = decode_3d_prediction(
+                            p3d[i],
+                            anchors_np[:, i],
+                            float(strides_np[i]),
+                            calib,
+                            img_w,
+                            img_h,
+                            face_3d_classes,
+                            complete_3d_classes,
+                            int(pred_cls_np[i]),
+                            pred_edge_60=pedge[i] if pedge is not None else None,
+                            bbox_xyxy=pred_bboxes[i],
+                        )
+                        if d is not None and d.get("corners_3d") is not None:
+                            draw_3d_box(
+                                im_pred_3d,
+                                d["corners_3d"],
+                                calib,
+                                d.get("face_center_2d"),
+                                d.get("face_color"),
+                                edge_points_2d=d.get("edge_points_2d"),
+                                edge_color=(0, 165, 255),
+                                thickness=line_thick,
+                            )
+
+        # --- Add labels to each quadrant ---
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        for im_q, label in [(im_gt_2d, "Target 2D"), (im_gt_3d, "Target 3D"),
+                            (im_pred_2d, "Pred 2D"), (im_pred_3d, "Pred 3D")]:
+            cv2.putText(im_q, label, (10, 40), font, 1.0, (0, 255, 255), 1, cv2.LINE_AA)
+
+        # --- Assemble 2x2 grid ---
+        top_row = np.hstack([im_gt_2d, im_gt_3d])
+        bot_row = np.hstack([im_pred_2d, im_pred_3d])
+        grid = np.vstack([top_row, bot_row])
+
+        header = f"{image_name or 'unknown'} | idx={global_sample_idx:04d} | {rank_tag} | mode={camera_mode}"
+        cv2.rectangle(grid, (0, 0), (grid.shape[1], 30), (32, 32, 32), -1)
+        cv2.putText(grid, header[:200], (10, 21), font, 0.55, (255, 255, 255), 1, cv2.LINE_AA)
+
+        # Convert BGR -> RGB for TensorBoard
+        grid_rgb = cv2.cvtColor(grid, cv2.COLOR_BGR2RGB)
+
+        WRITER.add_image(
+            f"val/{rank_tag}/{camera_mode}/sample{global_sample_idx:04d}_{image_tag}_2d3d",
+            grid_rgb,
+            dataformats="HWC",
+            global_step=epoch,
+        )
+
+    except Exception as e:
+        LOGGER.warning(f"{PREFIX}Failed to log 3D visualization: {e}")
+
+
+def on_val_end(validator) -> None:
+    """Clean up after validation."""
+    _reset_val_logging_state()
+
+
+def on_train_end(trainer) -> None:
+    """Flush and close the TensorBoard writer at the end of training."""
+    _reset_val_logging_state()
+    _close_writer()
+
+
+callbacks = (
+    {
+        "on_pretrain_routine_start": on_pretrain_routine_start,
+        "on_train_start": on_train_start,
+        "on_fit_epoch_end": on_fit_epoch_end,
+        "on_train_epoch_end": on_train_epoch_end,
+        "on_val_start": on_val_start,
+        "on_val_batch_end": on_val_batch_end,
+        "on_val_end": on_val_end,
+        "on_train_end": on_train_end,
+    }
+    if SummaryWriter
+    else {}
+)

ultralytics/utils/callbacks/wb.py

@@ -0,0 +1,193 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from ultralytics.utils import SETTINGS, TESTS_RUNNING
+from ultralytics.utils.torch_utils import model_info_for_loggers
+
+try:
+    assert not TESTS_RUNNING  # do not log pytest
+    assert SETTINGS["wandb"] is True  # verify integration is enabled
+    import wandb as wb
+
+    assert hasattr(wb, "__version__")  # verify package is not directory
+    _processed_plots = {}
+
+except (ImportError, AssertionError):
+    wb = None
+
+
+def _custom_table(x, y, classes, title="Precision Recall Curve", x_title="Recall", y_title="Precision"):
+    """Create and log a custom metric visualization table.
+
+    This function crafts a custom metric visualization that mimics the behavior of the default wandb precision-recall
+    curve while allowing for enhanced customization. The visual metric is useful for monitoring model performance across
+    different classes.
+
+    Args:
+        x (list): Values for the x-axis; expected to have length N.
+        y (list): Corresponding values for the y-axis; also expected to have length N.
+        classes (list): Labels identifying the class of each point; length N.
+        title (str, optional): Title for the plot.
+        x_title (str, optional): Label for the x-axis.
+        y_title (str, optional): Label for the y-axis.
+
+    Returns:
+        (wandb.Object): A wandb object suitable for logging, showcasing the crafted metric visualization.
+    """
+    import polars as pl  # scope for faster 'import ultralytics'
+    import polars.selectors as cs
+
+    df = pl.DataFrame({"class": classes, "y": y, "x": x}).with_columns(cs.numeric().round(3))
+    data = df.select(["class", "y", "x"]).rows()
+
+    fields = {"x": "x", "y": "y", "class": "class"}
+    string_fields = {"title": title, "x-axis-title": x_title, "y-axis-title": y_title}
+    return wb.plot_table(
+        "wandb/area-under-curve/v0",
+        wb.Table(data=data, columns=["class", "y", "x"]),
+        fields=fields,
+        string_fields=string_fields,
+    )
+
+
+def _plot_curve(
+    x,
+    y,
+    names=None,
+    id="precision-recall",
+    title="Precision Recall Curve",
+    x_title="Recall",
+    y_title="Precision",
+    num_x=100,
+    only_mean=False,
+):
+    """Log a metric curve visualization.
+
+    This function generates a metric curve based on input data and logs the visualization to wandb. The curve can
+    represent aggregated data (mean) or individual class data, depending on the 'only_mean' flag.
+
+    Args:
+        x (np.ndarray): Data points for the x-axis with length N.
+        y (np.ndarray): Corresponding data points for the y-axis with shape (C, N), where C is the number of classes.
+        names (list, optional): Names of the classes corresponding to the y-axis data; length C.
+        id (str, optional): Unique identifier for the logged data in wandb.
+        title (str, optional): Title for the visualization plot.
+        x_title (str, optional): Label for the x-axis.
+        y_title (str, optional): Label for the y-axis.
+        num_x (int, optional): Number of interpolated data points for visualization.
+        only_mean (bool, optional): Flag to indicate if only the mean curve should be plotted.
+
+    Notes:
+        The function leverages the '_custom_table' function to generate the actual visualization.
+    """
+    import numpy as np
+
+    # Create new x
+    if names is None:
+        names = []
+    x_new = np.linspace(x[0], x[-1], num_x).round(5)
+
+    # Create arrays for logging
+    x_log = x_new.tolist()
+    y_log = np.interp(x_new, x, np.mean(y, axis=0)).round(3).tolist()
+
+    if only_mean:
+        table = wb.Table(data=list(zip(x_log, y_log)), columns=[x_title, y_title])
+        wb.run.log({title: wb.plot.line(table, x_title, y_title, title=title)})
+    else:
+        classes = ["mean"] * len(x_log)
+        for i, yi in enumerate(y):
+            x_log.extend(x_new)  # add new x
+            y_log.extend(np.interp(x_new, x, yi))  # interpolate y to new x
+            classes.extend([names[i]] * len(x_new))  # add class names
+        wb.log({id: _custom_table(x_log, y_log, classes, title, x_title, y_title)}, commit=False)
+
+
+def _log_plots(plots, step):
+    """Log plots to WandB at a specific step if they haven't been logged already.
+
+    This function checks each plot in the input dictionary against previously processed plots and logs new or updated
+    plots to WandB at the specified step.
+
+    Args:
+        plots (dict): Dictionary of plots to log, where keys are plot names and values are dictionaries containing plot
+            metadata including timestamps.
+        step (int): The step/epoch at which to log the plots in the WandB run.
+
+    Notes:
+        The function uses a shallow copy of the plots dictionary to prevent modification during iteration.
+        Plots are identified by their stem name (filename without extension).
+        Each plot is logged as a WandB Image object.
+    """
+    for name, params in plots.copy().items():  # shallow copy to prevent plots dict changing during iteration
+        timestamp = params["timestamp"]
+        if _processed_plots.get(name) != timestamp:
+            wb.run.log({name.stem: wb.Image(str(name))}, step=step)
+            _processed_plots[name] = timestamp
+
+
+def on_pretrain_routine_start(trainer):
+    """Initialize and start wandb project if module is present."""
+    if not wb.run:
+        from datetime import datetime
+
+        name = str(trainer.args.name).replace("/", "-").replace(" ", "_")
+        wb.init(
+            project=str(trainer.args.project).replace("/", "-") if trainer.args.project else "Ultralytics",
+            name=name,
+            config=vars(trainer.args),
+            id=f"{name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",  # add unique id
+            dir=str(trainer.save_dir),
+        )
+
+
+def on_fit_epoch_end(trainer):
+    """Log training metrics and model information at the end of an epoch."""
+    _log_plots(trainer.plots, step=trainer.epoch + 1)
+    _log_plots(trainer.validator.plots, step=trainer.epoch + 1)
+    if trainer.epoch == 0:
+        wb.run.log(model_info_for_loggers(trainer), step=trainer.epoch + 1)
+    wb.run.log(trainer.metrics, step=trainer.epoch + 1, commit=True)  # commit forces sync
+
+
+def on_train_epoch_end(trainer):
+    """Log metrics and save images at the end of each training epoch."""
+    wb.run.log(trainer.label_loss_items(trainer.tloss, prefix="train"), step=trainer.epoch + 1)
+    wb.run.log(trainer.lr, step=trainer.epoch + 1)
+    if trainer.epoch == 1:
+        _log_plots(trainer.plots, step=trainer.epoch + 1)
+
+
+def on_train_end(trainer):
+    """Save the best model as an artifact and log final plots at the end of training."""
+    _log_plots(trainer.validator.plots, step=trainer.epoch + 1)
+    _log_plots(trainer.plots, step=trainer.epoch + 1)
+    art = wb.Artifact(type="model", name=f"run_{wb.run.id}_model")
+    if trainer.best.exists():
+        art.add_file(trainer.best)
+        wb.run.log_artifact(art, aliases=["best"])
+    # Check if we actually have plots to save
+    if trainer.args.plots and hasattr(trainer.validator.metrics, "curves_results"):
+        for curve_name, curve_values in zip(trainer.validator.metrics.curves, trainer.validator.metrics.curves_results):
+            x, y, x_title, y_title = curve_values
+            _plot_curve(
+                x,
+                y,
+                names=list(trainer.validator.metrics.names.values()),
+                id=f"curves/{curve_name}",
+                title=curve_name,
+                x_title=x_title,
+                y_title=y_title,
+            )
+    wb.run.finish()  # required or run continues on dashboard
+
+
+callbacks = (
+    {
+        "on_pretrain_routine_start": on_pretrain_routine_start,
+        "on_train_epoch_end": on_train_epoch_end,
+        "on_fit_epoch_end": on_fit_epoch_end,
+        "on_train_end": on_train_end,
+    }
+    if wb
+    else {}
+)

ultralytics/utils/cpu.py

@@ -0,0 +1,85 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import platform
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+
+class CPUInfo:
+    """Provide cross-platform CPU brand and model information.
+
+    Query platform-specific sources to retrieve a human-readable CPU descriptor and normalize it for consistent
+    presentation across macOS, Linux, and Windows. If platform-specific probing fails, generic platform identifiers are
+    used to ensure a stable string is always returned.
+
+    Methods:
+        name: Return the normalized CPU name using platform-specific sources with robust fallbacks.
+        _clean: Normalize and prettify common vendor brand strings and frequency patterns.
+        __str__: Return the normalized CPU name for string contexts.
+
+    Examples:
+        >>> CPUInfo.name()
+        'Apple M4 Pro'
+        >>> str(CPUInfo())
+        'Intel Core i7-9750H 2.60GHz'
+    """
+
+    @staticmethod
+    def name() -> str:
+        """Return a normalized CPU model string from platform-specific sources."""
+        try:
+            if sys.platform == "darwin":
+                # Query macOS sysctl for the CPU brand string
+                s = subprocess.run(
+                    ["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True
+                ).stdout.strip()
+                if s:
+                    return CPUInfo._clean(s)
+            elif sys.platform.startswith("linux"):
+                # Parse /proc/cpuinfo for the first "model name" entry
+                p = Path("/proc/cpuinfo")
+                if p.exists():
+                    for line in p.read_text(errors="ignore").splitlines():
+                        if "model name" in line:
+                            return CPUInfo._clean(line.split(":", 1)[1])
+            elif sys.platform.startswith("win"):
+                try:
+                    import winreg as wr
+
+                    with wr.OpenKey(wr.HKEY_LOCAL_MACHINE, r"HARDWARE\DESCRIPTION\System\CentralProcessor\0") as k:
+                        val, _ = wr.QueryValueEx(k, "ProcessorNameString")
+                        if val:
+                            return CPUInfo._clean(val)
+                except Exception:
+                    # Fall through to generic platform fallbacks on Windows registry access failure
+                    pass
+            # Generic platform fallbacks
+            s = platform.processor() or getattr(platform.uname(), "processor", "") or platform.machine()
+            return CPUInfo._clean(s or "Unknown CPU")
+        except Exception:
+            # Ensure a string is always returned even on unexpected failures
+            s = platform.processor() or platform.machine() or ""
+            return CPUInfo._clean(s or "Unknown CPU")
+
+    @staticmethod
+    def _clean(s: str) -> str:
+        """Normalize and prettify a raw CPU descriptor string."""
+        s = re.sub(r"\s+", " ", s.strip())
+        s = s.replace("(TM)", "").replace("(tm)", "").replace("(R)", "").replace("(r)", "").strip()
+        if m := re.search(r"(Intel.*?i\d[\w-]*) CPU @ ([\d.]+GHz)", s, re.I):
+            return f"{m.group(1)} {m.group(2)}"
+        if m := re.search(r"(AMD.*?Ryzen.*?[\w-]*) CPU @ ([\d.]+GHz)", s, re.I):
+            return f"{m.group(1)} {m.group(2)}"
+        return s
+
+    def __str__(self) -> str:
+        """Return the normalized CPU name."""
+        return self.name()
+
+
+if __name__ == "__main__":
+    print(CPUInfo.name())

ultralytics/utils/dist.py

@@ -0,0 +1,190 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import os
+import shutil
+import sys
+import tempfile
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING
+
+from . import USER_CONFIG_DIR
+from .torch_utils import TORCH_1_9
+
+if TYPE_CHECKING:
+    from ultralytics.engine.trainer import BaseTrainer
+
+
+_RUN_TIMESTAMP_STORES = {}
+
+
+def find_free_network_port() -> int:
+    """Find a free port on localhost.
+
+    It is useful in single-node training when we don't want to connect to a real main node but have to set the
+    `MASTER_PORT` environment variable.
+
+    Returns:
+        (int): The available network port number.
+    """
+    import socket
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]  # port
+
+
+def _get_run_sync_port(base_name: str) -> int:
+    """Derive a stable coordination port for sharing a run timestamp across nodes."""
+    override = os.getenv("ULTRALYTICS_RUN_SYNC_PORT")
+    if override:
+        try:
+            return int(override)
+        except ValueError as exc:
+            raise RuntimeError("ULTRALYTICS_RUN_SYNC_PORT must be an integer.") from exc
+
+    try:
+        master_port = int(os.getenv("MASTER_PORT", "29500"))
+    except ValueError:
+        master_port = 29500
+
+    offset = 11 + sum(ord(char) for char in base_name) % 97
+    return 1024 + ((master_port + offset - 1024) % 64512)
+
+
+def get_distributed_run_timestamp(base_name: str, world_size: int, rank: int, *, timeout: timedelta = timedelta(seconds=120)) -> str:
+    """Return a timestamp string that is shared across all distributed ranks."""
+    if world_size <= 1 or rank == -1:
+        return datetime.now().strftime("%Y%m%d%H%M%S")
+
+    master_addr = os.getenv("MASTER_ADDR")
+    if not master_addr:
+        raise RuntimeError("MASTER_ADDR is required to coordinate a distributed run name.")
+
+    from torch.distributed import TCPStore
+
+    sync_port = _get_run_sync_port(base_name)
+    safe_base_name = "".join(char if char.isalnum() else "_" for char in base_name)
+    key = f"ultralytics_run_timestamp_{safe_base_name}"
+    is_master = rank == 0
+    store_host = "0.0.0.0" if is_master else master_addr
+
+    try:
+        store = TCPStore(store_host, sync_port, world_size, is_master, timeout=timeout, wait_for_workers=True)
+    except TypeError:
+        store = TCPStore(store_host, sync_port, world_size, is_master, timeout)
+    except Exception as exc:
+        raise RuntimeError(
+            f"Failed to coordinate the distributed run timestamp via TCPStore at {master_addr}:{sync_port}. "
+            "Set ULTRALYTICS_RUN_SYNC_PORT to an open cross-node port if needed."
+        ) from exc
+
+    # Keep the TCPStore alive for the full process lifetime. Rank 0 hosts the store server, so if this object is
+    # garbage collected as soon as build_run_name() returns, the other ranks can lose the connection mid-handshake.
+    _RUN_TIMESTAMP_STORES[(master_addr, sync_port, world_size, rank)] = store
+
+    if is_master:
+        store.set(key, datetime.now().strftime("%Y%m%d%H%M%S"))
+
+    timestamp = store.get(key)
+    return timestamp.decode("utf-8") if isinstance(timestamp, bytes) else str(timestamp)
+
+
+def generate_ddp_file(trainer: BaseTrainer) -> str:
+    """Generate a DDP (Distributed Data Parallel) file for multi-GPU training.
+
+    This function creates a temporary Python file that enables distributed training across multiple GPUs. The file
+    contains the necessary configuration to initialize the trainer in a distributed environment.
+
+    Args:
+        trainer (ultralytics.engine.trainer.BaseTrainer): The trainer containing training configuration and arguments.
+            Must have args attribute and be a class instance.
+
+    Returns:
+        (str): Path to the generated temporary DDP file.
+
+    Notes:
+        The generated file is saved in the USER_CONFIG_DIR/DDP directory and includes:
+        - Trainer class import
+        - Configuration overrides from the trainer arguments
+        - Model path configuration
+        - Training initialization code
+    """
+    module, name = f"{trainer.__class__.__module__}.{trainer.__class__.__name__}".rsplit(".", 1)
+
+    content = f"""
+# Ultralytics Multi-GPU training temp file (should be automatically deleted after use)
+from pathlib import Path, PosixPath  # For model arguments stored as Path instead of str
+overrides = {vars(trainer.args)}
+
+if __name__ == "__main__":
+    from {module} import {name}
+    from ultralytics.utils import DEFAULT_CFG_DICT
+
+    cfg = DEFAULT_CFG_DICT.copy()
+    cfg.update(save_dir='')   # handle the extra key 'save_dir'
+    trainer = {name}(cfg=cfg, overrides=overrides)
+    trainer.args.model = "{getattr(trainer.hub_session, "model_url", trainer.args.model)}"
+    results = trainer.train()
+"""
+    (USER_CONFIG_DIR / "DDP").mkdir(exist_ok=True)
+    with tempfile.NamedTemporaryFile(
+        prefix="_temp_",
+        suffix=f"{id(trainer)}.py",
+        mode="w+",
+        encoding="utf-8",
+        dir=USER_CONFIG_DIR / "DDP",
+        delete=False,
+    ) as file:
+        file.write(content)
+    return file.name
+
+
+def generate_ddp_command(trainer: BaseTrainer) -> tuple[list[str], str]:
+    """Generate command for distributed training.
+
+    Args:
+        trainer (ultralytics.engine.trainer.BaseTrainer): The trainer containing configuration for distributed training.
+
+    Returns:
+        cmd (list[str]): The command to execute for distributed training.
+        file (str): Path to the temporary file created for DDP training.
+    """
+    import __main__  # noqa local import to avoid https://github.com/Lightning-AI/pytorch-lightning/issues/15218
+
+    if not trainer.resume:
+        shutil.rmtree(trainer.save_dir)  # remove the save_dir
+    file = generate_ddp_file(trainer)
+    dist_cmd = "torch.distributed.run" if TORCH_1_9 else "torch.distributed.launch"
+    port = find_free_network_port()
+    cmd = [
+        sys.executable,
+        "-m",
+        dist_cmd,
+        "--nproc_per_node",
+        f"{trainer.world_size}",
+        "--master_port",
+        f"{port}",
+        file,
+    ]
+    return cmd, file
+
+
+def ddp_cleanup(trainer: BaseTrainer, file: str) -> None:
+    """Delete temporary file if created during distributed data parallel (DDP) training.
+
+    This function checks if the provided file contains the trainer's ID in its name, indicating it was created as a
+    temporary file for DDP training, and deletes it if so.
+
+    Args:
+        trainer (ultralytics.engine.trainer.BaseTrainer): The trainer used for distributed training.
+        file (str): Path to the file that might need to be deleted.
+
+    Examples:
+        >>> trainer = YOLOTrainer()
+        >>> file = "/tmp/ddp_temp_123456789.py"
+        >>> ddp_cleanup(trainer, file)
+    """
+    if f"{id(trainer)}.py" in file:  # if temp_file suffix in file
+        os.remove(file)

ultralytics/utils/downloads.py

@@ -0,0 +1,535 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import re
+import shutil
+import subprocess
+from itertools import repeat
+from multiprocessing.pool import ThreadPool
+from pathlib import Path
+from urllib import parse, request
+
+from ultralytics.utils import ASSETS_URL, LOGGER, TQDM, checks, clean_url, emojis, is_online, url2file
+
+# Define Ultralytics GitHub assets maintained at https://github.com/ultralytics/assets
+GITHUB_ASSETS_REPO = "ultralytics/assets"
+GITHUB_ASSETS_NAMES = frozenset(
+    [f"yolov8{k}{suffix}.pt" for k in "nsmlx" for suffix in ("", "-cls", "-seg", "-pose", "-obb", "-oiv7")]
+    + [f"yolo11{k}{suffix}.pt" for k in "nsmlx" for suffix in ("", "-cls", "-seg", "-pose", "-obb")]
+    + [f"yolo12{k}{suffix}.pt" for k in "nsmlx" for suffix in ("",)]  # detect models only currently
+    + [f"yolo26{k}{suffix}.pt" for k in "nsmlx" for suffix in ("", "-cls", "-seg", "-pose", "-obb")]
+    + [f"yolov5{k}{resolution}u.pt" for k in "nsmlx" for resolution in ("", "6")]
+    + [f"yolov3{k}u.pt" for k in ("", "-spp", "-tiny")]
+    + [f"yolov8{k}-world.pt" for k in "smlx"]
+    + [f"yolov8{k}-worldv2.pt" for k in "smlx"]
+    + [f"yoloe-v8{k}{suffix}.pt" for k in "sml" for suffix in ("-seg", "-seg-pf")]
+    + [f"yoloe-11{k}{suffix}.pt" for k in "sml" for suffix in ("-seg", "-seg-pf")]
+    + [f"yoloe-26{k}{suffix}.pt" for k in "nsmlx" for suffix in ("-seg", "-seg-pf")]
+    + [f"yolov9{k}.pt" for k in "tsmce"]
+    + [f"yolov10{k}.pt" for k in "nsmblx"]
+    + [f"yolo_nas_{k}.pt" for k in "sml"]
+    + [f"sam_{k}.pt" for k in "bl"]
+    + [f"sam2_{k}.pt" for k in "blst"]
+    + [f"sam2.1_{k}.pt" for k in "blst"]
+    + [f"FastSAM-{k}.pt" for k in "sx"]
+    + [f"rtdetr-{k}.pt" for k in "lx"]
+    + [
+        "mobile_sam.pt",
+        "mobileclip_blt.ts",
+        "yolo11n-grayscale.pt",
+        "calibration_image_sample_data_20x128x128x3_float32.npy.zip",
+    ]
+)
+GITHUB_ASSETS_STEMS = frozenset(k.rpartition(".")[0] for k in GITHUB_ASSETS_NAMES)
+
+
+def is_url(url: str | Path, check: bool = False) -> bool:
+    """Validate if the given string is a URL and optionally check if the URL exists online.
+
+    Args:
+        url (str | Path): The string to be validated as a URL.
+        check (bool, optional): If True, performs an additional check to see if the URL exists online.
+
+    Returns:
+        (bool): True for a valid URL. If 'check' is True, also returns True if the URL exists online.
+
+    Examples:
+        >>> valid = is_url("https://www.example.com")
+        >>> valid_and_exists = is_url("https://www.example.com", check=True)
+    """
+    try:
+        url = str(url)
+        result = parse.urlparse(url)
+        if not (result.scheme and result.netloc):
+            return False
+        if check:
+            r = request.urlopen(request.Request(url, method="HEAD"), timeout=3)
+            return 200 <= r.getcode() < 400
+        return True
+    except Exception:
+        return False
+
+
+def delete_dsstore(path: str | Path, files_to_delete: tuple[str, ...] = (".DS_Store", "__MACOSX")) -> None:
+    """Delete all specified system files in a directory.
+
+    Args:
+        path (str | Path): The directory path where the files should be deleted.
+        files_to_delete (tuple[str, ...]): The files to be deleted.
+
+    Examples:
+        >>> from ultralytics.utils.downloads import delete_dsstore
+        >>> delete_dsstore("path/to/dir")
+
+    Notes:
+        ".DS_Store" files are created by the Apple operating system and contain metadata about folders and files. They
+        are hidden system files and can cause issues when transferring files between different operating systems.
+    """
+    for file in files_to_delete:
+        matches = list(Path(path).rglob(file))
+        LOGGER.info(f"Deleting {file} files: {matches}")
+        for f in matches:
+            f.unlink()
+
+
+def zip_directory(
+    directory: str | Path,
+    compress: bool = True,
+    exclude: tuple[str, ...] = (".DS_Store", "__MACOSX"),
+    progress: bool = True,
+) -> Path:
+    """Zip the contents of a directory, excluding specified files.
+
+    The resulting zip file is named after the directory and placed alongside it.
+
+    Args:
+        directory (str | Path): The path to the directory to be zipped.
+        compress (bool): Whether to compress the files while zipping.
+        exclude (tuple[str, ...], optional): A tuple of filename strings to be excluded.
+        progress (bool, optional): Whether to display a progress bar.
+
+    Returns:
+        (Path): The path to the resulting zip file.
+
+    Examples:
+        >>> from ultralytics.utils.downloads import zip_directory
+        >>> file = zip_directory("path/to/dir")
+    """
+    from zipfile import ZIP_DEFLATED, ZIP_STORED, ZipFile
+
+    delete_dsstore(directory)
+    directory = Path(directory)
+    if not directory.is_dir():
+        raise FileNotFoundError(f"Directory '{directory}' does not exist.")
+
+    # Zip with progress bar
+    files = [f for f in directory.rglob("*") if f.is_file() and all(x not in f.name for x in exclude)]  # files to zip
+    zip_file = directory.with_suffix(".zip")
+    compression = ZIP_DEFLATED if compress else ZIP_STORED
+    with ZipFile(zip_file, "w", compression) as f:
+        for file in TQDM(files, desc=f"Zipping {directory} to {zip_file}...", unit="files", disable=not progress):
+            f.write(file, file.relative_to(directory))
+
+    return zip_file  # return path to zip file
+
+
+def unzip_file(
+    file: str | Path,
+    path: str | Path | None = None,
+    exclude: tuple[str, ...] = (".DS_Store", "__MACOSX"),
+    exist_ok: bool = False,
+    progress: bool = True,
+) -> Path:
+    """Unzip a *.zip file to the specified path, excluding specified files.
+
+    If the zipfile does not contain a single top-level directory, the function will create a new directory with the same
+    name as the zipfile (without the extension) to extract its contents. If a path is not provided, the function will
+    use the parent directory of the zipfile as the default path.
+
+    Args:
+        file (str | Path): The path to the zipfile to be extracted.
+        path (str | Path, optional): The path to extract the zipfile to.
+        exclude (tuple[str, ...], optional): A tuple of filename strings to be excluded.
+        exist_ok (bool, optional): Whether to overwrite existing contents if they exist.
+        progress (bool, optional): Whether to display a progress bar.
+
+    Returns:
+        (Path): The path to the directory where the zipfile was extracted.
+
+    Raises:
+        BadZipFile: If the provided file does not exist or is not a valid zipfile.
+
+    Examples:
+        >>> from ultralytics.utils.downloads import unzip_file
+        >>> directory = unzip_file("path/to/file.zip")
+    """
+    from zipfile import BadZipFile, ZipFile, is_zipfile
+
+    if not (Path(file).exists() and is_zipfile(file)):
+        raise BadZipFile(f"File '{file}' does not exist or is a bad zip file.")
+    if path is None:
+        path = Path(file).parent  # default path
+
+    # Unzip the file contents
+    with ZipFile(file) as zipObj:
+        files = [f for f in zipObj.namelist() if all(x not in f for x in exclude)]
+        top_level_dirs = {Path(f).parts[0] for f in files}
+
+        # Decide to unzip directly or unzip into a directory
+        unzip_as_dir = len(top_level_dirs) == 1  # (len(files) > 1 and not files[0].endswith("/"))
+        if unzip_as_dir:
+            # Zip has 1 top-level directory
+            extract_path = path  # i.e. ../datasets
+            path = Path(path) / next(iter(top_level_dirs))  # i.e. extract coco8/ dir to ../datasets/
+        else:
+            # Zip has multiple files at top level
+            path = extract_path = Path(path) / Path(file).stem  # i.e. extract multiple files to ../datasets/coco8/
+
+        # Check if destination directory already exists and contains files
+        if path.exists() and any(path.iterdir()) and not exist_ok:
+            # If it exists and is not empty, return the path without unzipping
+            LOGGER.warning(f"Skipping {file} unzip as destination directory {path} is not empty.")
+            return path
+
+        for f in TQDM(files, desc=f"Unzipping {file} to {Path(path).resolve()}...", unit="files", disable=not progress):
+            # Ensure the file is within the extract_path to avoid path traversal security vulnerability
+            if ".." in Path(f).parts:
+                LOGGER.warning(f"Potentially insecure file path: {f}, skipping extraction.")
+                continue
+            zipObj.extract(f, extract_path)
+
+    return path  # return unzip dir
+
+
+def check_disk_space(
+    file_bytes: int,
+    path: str | Path = Path.cwd(),
+    sf: float = 1.5,
+    hard: bool = True,
+) -> bool:
+    """Check if there is sufficient disk space to download and store a file.
+
+    Args:
+        file_bytes (int): The file size in bytes.
+        path (str | Path, optional): The path or drive to check the available free space on.
+        sf (float, optional): Safety factor, the multiplier for the required free space.
+        hard (bool, optional): Whether to throw an error or not on insufficient disk space.
+
+    Returns:
+        (bool): True if there is sufficient disk space, False otherwise.
+    """
+    _total, _used, free = shutil.disk_usage(path)  # bytes
+    if file_bytes * sf < free:
+        return True  # sufficient space
+
+    def fmt_bytes(b):
+        return f"{b / (1 << 20):.1f} MB" if b < (1 << 30) else f"{b / (1 << 30):.3f} GB"
+
+    # Insufficient space
+    text = (
+        f"Insufficient free disk space {fmt_bytes(free)} < {fmt_bytes(int(file_bytes * sf))} required, "
+        f"Please free {fmt_bytes(int(file_bytes * sf - free))} additional disk space and try again."
+    )
+    if hard:
+        raise MemoryError(text)
+    LOGGER.warning(text)
+    return False
+
+
+def get_google_drive_file_info(link: str) -> tuple[str, str | None]:
+    """Retrieve the direct download link and filename for a shareable Google Drive file link.
+
+    Args:
+        link (str): The shareable link of the Google Drive file.
+
+    Returns:
+        url (str): Direct download URL for the Google Drive file.
+        filename (str | None): Original filename of the Google Drive file. If filename extraction fails, returns None.
+
+    Examples:
+        >>> from ultralytics.utils.downloads import get_google_drive_file_info
+        >>> link = "https://drive.google.com/file/d/1cqT-cJgANNrhIHCrEufUYhQ4RqiWG_lJ/view?usp=drive_link"
+        >>> url, filename = get_google_drive_file_info(link)
+    """
+    import requests  # scoped as slow import
+
+    file_id = link.split("/d/")[1].split("/view", 1)[0]
+    drive_url = f"https://drive.google.com/uc?export=download&id={file_id}"
+    filename = None
+
+    # Start session
+    with requests.Session() as session:
+        response = session.get(drive_url, stream=True)
+        if "quota exceeded" in str(response.content.lower()):
+            raise ConnectionError(
+                emojis(
+                    f"❌  Google Drive file download quota exceeded. "
+                    f"Please try again later or download this file manually at {link}."
+                )
+            )
+        for k, v in response.cookies.items():
+            if k.startswith("download_warning"):
+                drive_url += f"&confirm={v}"  # v is token
+        if cd := response.headers.get("content-disposition"):
+            filename = re.findall('filename="(.+)"', cd)[0]
+    return drive_url, filename
+
+
+def safe_download(
+    url: str | Path,
+    file: str | Path | None = None,
+    dir: str | Path | None = None,
+    unzip: bool = True,
+    delete: bool = False,
+    curl: bool = False,
+    retry: int = 3,
+    min_bytes: float = 1e0,
+    exist_ok: bool = False,
+    progress: bool = True,
+) -> Path | str:
+    """Download files from a URL with options for retrying, unzipping, and deleting the downloaded file. Enhanced with
+    robust partial download detection using Content-Length validation.
+
+    Args:
+        url (str | Path): The URL of the file to be downloaded.
+        file (str | Path, optional): The filename of the downloaded file. If not provided, the file will be saved with
+            the same name as the URL.
+        dir (str | Path, optional): The directory to save the downloaded file. If not provided, the file will be saved
+            in the current working directory.
+        unzip (bool, optional): Whether to unzip the downloaded file.
+        delete (bool, optional): Whether to delete the downloaded file after unzipping.
+        curl (bool, optional): Whether to use curl command line tool for downloading.
+        retry (int, optional): The number of times to retry the download in case of failure.
+        min_bytes (float, optional): The minimum number of bytes that the downloaded file should have, to be considered
+            a successful download.
+        exist_ok (bool, optional): Whether to overwrite existing contents during unzipping.
+        progress (bool, optional): Whether to display a progress bar during the download.
+
+    Returns:
+        (Path | str): The path to the downloaded file or extracted directory.
+
+    Examples:
+        >>> from ultralytics.utils.downloads import safe_download
+        >>> link = "https://ultralytics.com/assets/bus.jpg"
+        >>> path = safe_download(link)
+    """
+    gdrive = url.startswith("https://drive.google.com/")  # check if the URL is a Google Drive link
+    if gdrive:
+        url, file = get_google_drive_file_info(url)
+    url = url.replace(" ", "%20")  # encode spaces for curl/urllib compatibility
+
+    f = Path(dir or ".") / (file or url2file(url))  # URL converted to filename
+    if "://" not in str(url) and Path(url).is_file():  # URL exists ('://' check required in Windows Python<3.10)
+        f = Path(url)  # filename
+    elif not f.is_file():  # URL and file do not exist
+        uri = (url if gdrive else clean_url(url)).replace(ASSETS_URL, "https://ultralytics.com/assets")  # clean
+        desc = f"Downloading {uri} to '{f}'"
+        f.parent.mkdir(parents=True, exist_ok=True)  # make directory if missing
+        curl_installed = shutil.which("curl")
+        for i in range(retry + 1):
+            try:
+                if (curl or i > 0) and curl_installed:  # curl download with retry, continue
+                    s = "sS" * (not progress)  # silent
+                    r = subprocess.run(["curl", "-#", f"-{s}L", url, "-o", f, "--retry", "3", "-C", "-"]).returncode
+                    assert r == 0, f"Curl return value {r}"
+                    expected_size = None  # Can't get size with curl
+                else:  # urllib download
+                    with request.urlopen(url) as response:
+                        expected_size = int(response.getheader("Content-Length", 0))
+                        if i == 0 and expected_size > 1048576:
+                            check_disk_space(expected_size, path=f.parent)
+                        buffer_size = max(8192, min(1048576, expected_size // 1000)) if expected_size else 8192
+                        with TQDM(
+                            total=expected_size,
+                            desc=desc,
+                            disable=not progress,
+                            unit="B",
+                            unit_scale=True,
+                            unit_divisor=1024,
+                        ) as pbar:
+                            with open(f, "wb") as f_opened:
+                                while True:
+                                    data = response.read(buffer_size)
+                                    if not data:
+                                        break
+                                    f_opened.write(data)
+                                    pbar.update(len(data))
+
+                if f.exists():
+                    file_size = f.stat().st_size
+                    if file_size > min_bytes:
+                        # Check if download is complete (only if we have expected_size)
+                        if expected_size and file_size != expected_size:
+                            LOGGER.warning(
+                                f"Partial download: {file_size}/{expected_size} bytes ({file_size / expected_size * 100:.1f}%)"
+                            )
+                        else:
+                            break  # success
+                    f.unlink()  # remove partial downloads
+            except MemoryError:
+                raise  # Re-raise immediately - no point retrying if insufficient disk space
+            except Exception as e:
+                if i == 0 and not is_online():
+                    raise ConnectionError(emojis(f"❌  Download failure for {uri}. Environment may be offline.")) from e
+                elif i >= retry:
+                    raise ConnectionError(emojis(f"❌  Download failure for {uri}. Retry limit reached. {e}")) from e
+                LOGGER.warning(f"Download failure, retrying {i + 1}/{retry} {uri}... {e}")
+
+    if unzip and f.exists() and f.suffix in {"", ".zip", ".tar", ".gz"}:
+        from zipfile import is_zipfile
+
+        unzip_dir = (dir or f.parent).resolve()  # unzip to dir if provided else unzip in place
+        if is_zipfile(f):
+            unzip_dir = unzip_file(file=f, path=unzip_dir, exist_ok=exist_ok, progress=progress)  # unzip
+        elif f.suffix in {".tar", ".gz"}:
+            LOGGER.info(f"Unzipping {f} to {unzip_dir}...")
+            subprocess.run(["tar", "xf" if f.suffix == ".tar" else "xfz", f, "--directory", unzip_dir], check=True)
+        if delete:
+            f.unlink()  # remove zip
+        return unzip_dir
+    return f
+
+
+def get_github_assets(
+    repo: str = "ultralytics/assets",
+    version: str = "latest",
+    retry: bool = False,
+) -> tuple[str, list[str]]:
+    """Retrieve the specified version's tag and assets from a GitHub repository.
+
+    If the version is not specified, the function fetches the latest release assets.
+
+    Args:
+        repo (str, optional): The GitHub repository in the format 'owner/repo'.
+        version (str, optional): The release version to fetch assets from.
+        retry (bool, optional): Flag to retry the request in case of a failure.
+
+    Returns:
+        tag (str): The release tag.
+        assets (list[str]): A list of asset names.
+
+    Examples:
+        >>> tag, assets = get_github_assets(repo="ultralytics/assets", version="latest")
+    """
+    import requests  # scoped as slow import
+
+    if version != "latest":
+        version = f"tags/{version}"  # i.e. tags/v6.2
+    url = f"https://api.github.com/repos/{repo}/releases/{version}"
+    r = requests.get(url)  # github api
+    if r.status_code != 200 and r.reason != "rate limit exceeded" and retry:  # failed and not 403 rate limit exceeded
+        r = requests.get(url)  # try again
+    if r.status_code != 200:
+        LOGGER.warning(f"GitHub assets check failure for {url}: {r.status_code} {r.reason}")
+        return "", []
+    data = r.json()
+    return data["tag_name"], [x["name"] for x in data["assets"]]  # tag, assets i.e. ['yolo26n.pt', 'yolo11s.pt', ...]
+
+
+def attempt_download_asset(
+    file: str | Path,
+    repo: str = "ultralytics/assets",
+    release: str = "v8.4.0",
+    **kwargs,
+) -> str:
+    """Attempt to download a file from GitHub release assets if it is not found locally.
+
+    Args:
+        file (str | Path): The filename or file path to be downloaded.
+        repo (str, optional): The GitHub repository in the format 'owner/repo'.
+        release (str, optional): The specific release version to be downloaded.
+        **kwargs (Any): Additional keyword arguments for the download process.
+
+    Returns:
+        (str): The path to the downloaded file.
+
+    Examples:
+        >>> file_path = attempt_download_asset("yolo26n.pt", repo="ultralytics/assets", release="latest")
+    """
+    from ultralytics.utils import SETTINGS  # scoped for circular import
+
+    # YOLOv3/5u updates
+    file = str(file)
+    file = checks.check_yolov5u_filename(file)
+    file = Path(file.strip().replace("'", ""))
+    if file.exists():
+        return str(file)
+    elif (SETTINGS["weights_dir"] / file).exists():
+        return str(SETTINGS["weights_dir"] / file)
+    else:
+        # URL specified
+        name = Path(parse.unquote(str(file))).name  # decode '%2F' to '/' etc.
+        download_url = f"https://github.com/{repo}/releases/download"
+        if str(file).startswith(("http:/", "https:/")):  # download
+            url = str(file).replace(":/", "://")  # Pathlib turns :// -> :/
+            file = url2file(name)  # parse authentication https://url.com/file.txt?auth...
+            if Path(file).is_file():
+                LOGGER.info(f"Found {clean_url(url)} locally at {file}")  # file already exists
+            else:
+                safe_download(url=url, file=file, min_bytes=1e5, **kwargs)
+
+        elif repo == GITHUB_ASSETS_REPO and name in GITHUB_ASSETS_NAMES:
+            safe_download(url=f"{download_url}/{release}/{name}", file=file, min_bytes=1e5, **kwargs)
+
+        else:
+            tag, assets = get_github_assets(repo, release)
+            if not assets:
+                tag, assets = get_github_assets(repo)  # latest release
+            if name in assets:
+                safe_download(url=f"{download_url}/{tag}/{name}", file=file, min_bytes=1e5, **kwargs)
+
+        return str(file)
+
+
+def download(
+    url: str | list[str] | Path,
+    dir: Path = Path.cwd(),
+    unzip: bool = True,
+    delete: bool = False,
+    curl: bool = False,
+    threads: int = 1,
+    retry: int = 3,
+    exist_ok: bool = False,
+) -> None:
+    """Download files from specified URLs to a given directory.
+
+    Supports concurrent downloads if multiple threads are specified.
+
+    Args:
+        url (str | list[str] | Path): The URL or list of URLs of the files to be downloaded.
+        dir (Path, optional): The directory where the files will be saved.
+        unzip (bool, optional): Flag to unzip the files after downloading.
+        delete (bool, optional): Flag to delete the zip files after extraction.
+        curl (bool, optional): Flag to use curl for downloading.
+        threads (int, optional): Number of threads to use for concurrent downloads.
+        retry (int, optional): Number of retries in case of download failure.
+        exist_ok (bool, optional): Whether to overwrite existing contents during unzipping.
+
+    Examples:
+        >>> download("https://ultralytics.com/assets/example.zip", dir="path/to/dir", unzip=True)
+    """
+    dir = Path(dir)
+    dir.mkdir(parents=True, exist_ok=True)  # make directory
+    urls = [url] if isinstance(url, (str, Path)) else url
+    if threads > 1:
+        LOGGER.info(f"Downloading {len(urls)} file(s) with {threads} threads to {dir}...")
+        with ThreadPool(threads) as pool:
+            pool.map(
+                lambda x: safe_download(
+                    url=x[0],
+                    dir=x[1],
+                    unzip=unzip,
+                    delete=delete,
+                    curl=curl,
+                    retry=retry,
+                    exist_ok=exist_ok,
+                    progress=True,
+                ),
+                zip(urls, repeat(dir)),
+            )
+            pool.close()
+            pool.join()
+    else:
+        for u in urls:
+            safe_download(url=u, dir=dir, unzip=unzip, delete=delete, curl=curl, retry=retry, exist_ok=exist_ok)

ultralytics/utils/errors.py

@@ -0,0 +1,35 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from ultralytics.utils import emojis
+
+
+class HUBModelError(Exception):
+    """Exception raised when a model cannot be found or retrieved from Ultralytics HUB.
+
+    This custom exception is used specifically for handling errors related to model fetching in Ultralytics YOLO. The
+    error message is processed to include emojis for better user experience.
+
+    Attributes:
+        message (str): The error message displayed when the exception is raised.
+
+    Methods:
+        __init__: Initialize the HUBModelError with a custom message.
+
+    Examples:
+        >>> try:
+        ...     # Code that might fail to find a model
+        ...     raise HUBModelError("Custom model not found message")
+        ... except HUBModelError as e:
+        ...     print(e)  # Displays the emoji-enhanced error message
+    """
+
+    def __init__(self, message: str = "Model not found. Please check model URL and try again."):
+        """Initialize a HUBModelError exception.
+
+        This exception is raised when a requested model is not found or cannot be retrieved from Ultralytics HUB. The
+        message is processed to include emojis for better user experience.
+
+        Args:
+            message (str, optional): The error message to display when the exception is raised.
+        """
+        super().__init__(emojis(message))

ultralytics/utils/events.py

@@ -0,0 +1,113 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import json
+import random
+import time
+from pathlib import Path
+from threading import Thread
+from urllib.request import Request, urlopen
+
+from ultralytics import SETTINGS, __version__
+from ultralytics.utils import ARGV, ENVIRONMENT, GIT, IS_PIP_PACKAGE, ONLINE, PYTHON_VERSION, RANK, TESTS_RUNNING
+from ultralytics.utils.downloads import GITHUB_ASSETS_NAMES
+from ultralytics.utils.torch_utils import get_cpu_info
+
+
+def _post(url: str, data: dict, timeout: float = 5.0) -> None:
+    """Send a one-shot JSON POST request."""
+    try:
+        body = json.dumps(data, separators=(",", ":")).encode()  # compact JSON
+        req = Request(url, data=body, headers={"Content-Type": "application/json"})
+        urlopen(req, timeout=timeout).close()
+    except Exception:
+        pass
+
+
+class Events:
+    """Collect and send anonymous usage analytics with rate-limiting.
+
+    Event collection and transmission are enabled when sync is enabled in settings, the current process is rank -1 or 0,
+    tests are not running, the environment is online, and the installation source is either pip or the official
+    Ultralytics GitHub repository.
+
+    Attributes:
+        url (str): Measurement Protocol endpoint for receiving anonymous events.
+        events (list[dict]): In-memory queue of event payloads awaiting transmission.
+        rate_limit (float): Minimum time in seconds between POST requests.
+        t (float): Timestamp of the last transmission in seconds since the epoch.
+        metadata (dict): Static metadata describing runtime, installation source, and environment.
+        enabled (bool): Flag indicating whether analytics collection is active.
+
+    Methods:
+        __init__: Initialize the event queue, rate limiter, and runtime metadata.
+        __call__: Queue an event and trigger a non-blocking send when the rate limit elapses.
+    """
+
+    url = "https://www.google-analytics.com/mp/collect?measurement_id=G-X8NCJYTQXM&api_secret=QLQrATrNSwGRFRLE-cbHJw"
+
+    def __init__(self) -> None:
+        """Initialize the Events instance with queue, rate limiter, and environment metadata."""
+        self.events = []  # pending events
+        self.rate_limit = 30.0  # rate limit (seconds)
+        self.t = 0.0  # last send timestamp (seconds)
+        self.metadata = {
+            "cli": Path(ARGV[0]).name == "yolo",
+            "install": "git" if GIT.is_repo else "pip" if IS_PIP_PACKAGE else "other",
+            "python": PYTHON_VERSION.rsplit(".", 1)[0],  # i.e. 3.13
+            "CPU": get_cpu_info(),
+            # "GPU": get_gpu_info(index=0) if cuda else None,
+            "version": __version__,
+            "env": ENVIRONMENT,
+            "session_id": round(random.random() * 1e15),
+            "engagement_time_msec": 1000,
+        }
+        self.enabled = (
+            SETTINGS["sync"]
+            and RANK in {-1, 0}
+            and not TESTS_RUNNING
+            and ONLINE
+            and (IS_PIP_PACKAGE or GIT.origin == "https://github.com/ultralytics/ultralytics.git")
+        )
+
+    def __call__(self, cfg, device=None) -> None:
+        """Queue an event and flush the queue asynchronously when the rate limit elapses.
+
+        Args:
+            cfg (IterableSimpleNamespace): The configuration object containing mode and task information.
+            device (torch.device | str, optional): The device type (e.g., 'cpu', 'cuda').
+        """
+        if not self.enabled:
+            # Events disabled, do nothing
+            return
+
+        # Attempt to enqueue a new event
+        if len(self.events) < 25:  # Queue limited to 25 events to bound memory and traffic
+            params = {
+                **self.metadata,
+                "task": cfg.task,
+                "model": cfg.model if cfg.model in GITHUB_ASSETS_NAMES else "custom",
+                "device": str(device),
+            }
+            if cfg.mode == "export":
+                params["format"] = cfg.format
+            self.events.append({"name": cfg.mode, "params": params})
+
+        # Check rate limit and return early if under limit
+        t = time.time()
+        if (t - self.t) < self.rate_limit:
+            return
+
+        # Overrate limit: send a snapshot of queued events in a background thread
+        payload_events = list(self.events)  # snapshot to avoid race with queue reset
+        Thread(
+            target=_post,
+            args=(self.url, {"client_id": SETTINGS["uuid"], "events": payload_events}),  # SHA-256 anonymized
+            daemon=True,
+        ).start()
+
+        # Reset queue and rate limit timer
+        self.events = []
+        self.t = t
+
+
+events = Events()

ultralytics/utils/export/__init__.py

@@ -0,0 +1,17 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from .engine import onnx2engine, torch2onnx
+from .executorch import torch2executorch
+from .imx import torch2imx
+from .tensorflow import keras2pb, onnx2saved_model, pb2tfjs, tflite2edgetpu
+
+__all__ = [
+    "keras2pb",
+    "onnx2engine",
+    "onnx2saved_model",
+    "pb2tfjs",
+    "tflite2edgetpu",
+    "torch2executorch",
+    "torch2imx",
+    "torch2onnx",
+]

ultralytics/utils/export/engine.py

@@ -0,0 +1,246 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+from ultralytics.utils import IS_JETSON, LOGGER
+from ultralytics.utils.torch_utils import TORCH_2_4
+
+
+def torch2onnx(
+    torch_model: torch.nn.Module,
+    im: torch.Tensor,
+    onnx_file: str,
+    opset: int = 14,
+    input_names: list[str] = ["images"],
+    output_names: list[str] = ["output0"],
+    dynamic: bool | dict = False,
+) -> None:
+    """Export a PyTorch model to ONNX format.
+
+    Args:
+        torch_model (torch.nn.Module): The PyTorch model to export.
+        im (torch.Tensor): Example input tensor for the model.
+        onnx_file (str): Path to save the exported ONNX file.
+        opset (int): ONNX opset version to use for export.
+        input_names (list[str]): List of input tensor names.
+        output_names (list[str]): List of output tensor names.
+        dynamic (bool | dict, optional): Whether to enable dynamic axes.
+
+    Notes:
+        Setting `do_constant_folding=True` may cause issues with DNN inference for torch>=1.12.
+    """
+    kwargs = {"dynamo": False} if TORCH_2_4 else {}
+    torch.onnx.export(
+        torch_model,
+        im,
+        onnx_file,
+        verbose=False,
+        opset_version=opset,
+        do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
+        input_names=input_names,
+        output_names=output_names,
+        dynamic_axes=dynamic or None,
+        **kwargs,
+    )
+
+
+def onnx2engine(
+    onnx_file: str,
+    engine_file: str | None = None,
+    workspace: int | None = None,
+    half: bool = False,
+    int8: bool = False,
+    dynamic: bool = False,
+    shape: tuple[int, int, int, int] = (1, 3, 640, 640),
+    dla: int | None = None,
+    dataset=None,
+    metadata: dict | None = None,
+    verbose: bool = False,
+    prefix: str = "",
+) -> None:
+    """Export a YOLO model to TensorRT engine format.
+
+    Args:
+        onnx_file (str): Path to the ONNX file to be converted.
+        engine_file (str | None): Path to save the generated TensorRT engine file.
+        workspace (int | None): Workspace size in GB for TensorRT.
+        half (bool, optional): Enable FP16 precision.
+        int8 (bool, optional): Enable INT8 precision.
+        dynamic (bool, optional): Enable dynamic input shapes.
+        shape (tuple[int, int, int, int], optional): Input shape (batch, channels, height, width).
+        dla (int | None): DLA core to use (Jetson devices only).
+        dataset (ultralytics.data.build.InfiniteDataLoader, optional): Dataset for INT8 calibration.
+        metadata (dict | None): Metadata to include in the engine file.
+        verbose (bool, optional): Enable verbose logging.
+        prefix (str, optional): Prefix for log messages.
+
+    Raises:
+        ValueError: If DLA is enabled on non-Jetson devices or required precision is not set.
+        RuntimeError: If the ONNX file cannot be parsed.
+
+    Notes:
+        TensorRT version compatibility is handled for workspace size and engine building.
+        INT8 calibration requires a dataset and generates a calibration cache.
+        Metadata is serialized and written to the engine file if provided.
+    """
+    import tensorrt as trt
+
+    engine_file = engine_file or Path(onnx_file).with_suffix(".engine")
+
+    logger = trt.Logger(trt.Logger.INFO)
+    if verbose:
+        logger.min_severity = trt.Logger.Severity.VERBOSE
+
+    # Engine builder
+    builder = trt.Builder(logger)
+    config = builder.create_builder_config()
+    workspace_bytes = int((workspace or 0) * (1 << 30))
+    is_trt10 = int(trt.__version__.split(".", 1)[0]) >= 10  # is TensorRT >= 10
+    if is_trt10 and workspace_bytes > 0:
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)
+    elif workspace_bytes > 0:  # TensorRT versions 7, 8
+        config.max_workspace_size = workspace_bytes
+    flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+    network = builder.create_network(flag)
+    half = builder.platform_has_fast_fp16 and half
+    int8 = builder.platform_has_fast_int8 and int8
+
+    # Optionally switch to DLA if enabled
+    if dla is not None:
+        if not IS_JETSON:
+            raise ValueError("DLA is only available on NVIDIA Jetson devices")
+        LOGGER.info(f"{prefix} enabling DLA on core {dla}...")
+        if not half and not int8:
+            raise ValueError(
+                "DLA requires either 'half=True' (FP16) or 'int8=True' (INT8) to be enabled. Please enable one of them and try again."
+            )
+        config.default_device_type = trt.DeviceType.DLA
+        config.DLA_core = int(dla)
+        config.set_flag(trt.BuilderFlag.GPU_FALLBACK)
+
+    # Read ONNX file
+    parser = trt.OnnxParser(network, logger)
+    if not parser.parse_from_file(onnx_file):
+        raise RuntimeError(f"failed to load ONNX file: {onnx_file}")
+
+    # Network inputs
+    inputs = [network.get_input(i) for i in range(network.num_inputs)]
+    outputs = [network.get_output(i) for i in range(network.num_outputs)]
+    for inp in inputs:
+        LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')
+    for out in outputs:
+        LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
+
+    if dynamic:
+        profile = builder.create_optimization_profile()
+        min_shape = (1, shape[1], 32, 32)  # minimum input shape
+        max_shape = (*shape[:2], *(int(max(2, workspace or 2) * d) for d in shape[2:]))  # max input shape
+        for inp in inputs:
+            profile.set_shape(inp.name, min=min_shape, opt=shape, max=max_shape)
+        config.add_optimization_profile(profile)
+        if int8 and not is_trt10:  # deprecated in TensorRT 10, causes internal errors
+            config.set_calibration_profile(profile)
+
+    LOGGER.info(f"{prefix} building {'INT8' if int8 else 'FP' + ('16' if half else '32')} engine as {engine_file}")
+    if int8:
+        config.set_flag(trt.BuilderFlag.INT8)
+        config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
+
+        class EngineCalibrator(trt.IInt8Calibrator):
+            """Custom INT8 calibrator for TensorRT engine optimization.
+
+            This calibrator provides the necessary interface for TensorRT to perform INT8 quantization calibration using
+            a dataset. It handles batch generation, caching, and calibration algorithm selection.
+
+            Attributes:
+                dataset: Dataset for calibration.
+                data_iter: Iterator over the calibration dataset.
+                algo (trt.CalibrationAlgoType): Calibration algorithm type.
+                batch (int): Batch size for calibration.
+                cache (Path): Path to save the calibration cache.
+
+            Methods:
+                get_algorithm: Get the calibration algorithm to use.
+                get_batch_size: Get the batch size to use for calibration.
+                get_batch: Get the next batch to use for calibration.
+                read_calibration_cache: Use existing cache instead of calibrating again.
+                write_calibration_cache: Write calibration cache to disk.
+            """
+
+            def __init__(
+                self,
+                dataset,  # ultralytics.data.build.InfiniteDataLoader
+                cache: str = "",
+            ) -> None:
+                """Initialize the INT8 calibrator with dataset and cache path."""
+                trt.IInt8Calibrator.__init__(self)
+                self.dataset = dataset
+                self.data_iter = iter(dataset)
+                self.algo = (
+                    trt.CalibrationAlgoType.ENTROPY_CALIBRATION_2  # DLA quantization needs ENTROPY_CALIBRATION_2
+                    if dla is not None
+                    else trt.CalibrationAlgoType.MINMAX_CALIBRATION
+                )
+                self.batch = dataset.batch_size
+                self.cache = Path(cache)
+
+            def get_algorithm(self) -> trt.CalibrationAlgoType:
+                """Get the calibration algorithm to use."""
+                return self.algo
+
+            def get_batch_size(self) -> int:
+                """Get the batch size to use for calibration."""
+                return self.batch or 1
+
+            def get_batch(self, names) -> list[int] | None:
+                """Get the next batch to use for calibration, as a list of device memory pointers."""
+                try:
+                    im0s = next(self.data_iter)["img"] / 255.0
+                    im0s = im0s.to("cuda") if im0s.device.type == "cpu" else im0s
+                    return [int(im0s.data_ptr())]
+                except StopIteration:
+                    # Return None to signal to TensorRT there is no calibration data remaining
+                    return None
+
+            def read_calibration_cache(self) -> bytes | None:
+                """Use existing cache instead of calibrating again, otherwise, implicitly return None."""
+                if self.cache.exists() and self.cache.suffix == ".cache":
+                    return self.cache.read_bytes()
+
+            def write_calibration_cache(self, cache: bytes) -> None:
+                """Write calibration cache to disk."""
+                _ = self.cache.write_bytes(cache)
+
+        # Load dataset w/ builder (for batching) and calibrate
+        config.int8_calibrator = EngineCalibrator(
+            dataset=dataset,
+            cache=str(Path(onnx_file).with_suffix(".cache")),
+        )
+
+    elif half:
+        config.set_flag(trt.BuilderFlag.FP16)
+
+    # Write file
+    if is_trt10:
+        # TensorRT 10+ returns bytes directly, not a context manager
+        engine = builder.build_serialized_network(network, config)
+        if engine is None:
+            raise RuntimeError("TensorRT engine build failed, check logs for errors")
+        with open(engine_file, "wb") as t:
+            if metadata is not None:
+                meta = json.dumps(metadata)
+                t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
+                t.write(meta.encode())
+            t.write(engine)
+    else:
+        with builder.build_engine(network, config) as engine, open(engine_file, "wb") as t:
+            if metadata is not None:
+                meta = json.dumps(metadata)
+                t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
+                t.write(meta.encode())
+            t.write(engine.serialize())

ultralytics/utils/export/executorch.py

@@ -0,0 +1,79 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from functools import partial
+from pathlib import Path
+
+import torch
+
+from ultralytics.nn.modules import Pose, Pose26
+from ultralytics.utils import LOGGER, YAML
+
+
+def executorch_wrapper(model: torch.nn.Module) -> torch.nn.Module:
+    """Apply ExecuTorch-specific model patches required for export/runtime compatibility."""
+    import types
+
+    for m in model.modules():
+        if not isinstance(m, Pose):
+            continue
+        m.kpts_decode = types.MethodType(partial(_executorch_kpts_decode, is_pose26=type(m) is Pose26), m)
+    return model
+
+
+def _executorch_kpts_decode(self, kpts: torch.Tensor, is_pose26: bool = False) -> torch.Tensor:
+    """Decode pose keypoints for ExecuTorch export with XNNPACK-safe broadcasting."""
+    ndim = self.kpt_shape[1]
+    bs = kpts.shape[0]
+    y = kpts.view(bs, *self.kpt_shape, -1)
+
+    # XNNPACK requires explicit dim matching for broadcasting, expand 2D tensors to 4D.
+    anchors = self.anchors[None, None]
+    strides = self.strides[None, None]
+    a = ((y[:, :, :2] + anchors) if is_pose26 else (y[:, :, :2] * 2.0 + (anchors - 0.5))) * strides
+    if ndim == 3:
+        a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
+    return a.view(bs, self.nk, -1)
+
+
+def torch2executorch(
+    model: torch.nn.Module,
+    file: Path | str,
+    sample_input: torch.Tensor,
+    metadata: dict | None = None,
+    prefix: str = "",
+) -> str:
+    """Export a PyTorch model to ExecuTorch format.
+
+    Args:
+        model (torch.nn.Module): The PyTorch model to export.
+        file (Path | str): Source model file path used to derive output names.
+        sample_input (torch.Tensor): Example input tensor for tracing/export.
+        metadata (dict | None, optional): Optional metadata to save as YAML.
+        prefix (str, optional): Prefix for log messages.
+
+    Returns:
+        (str): Path to the exported ExecuTorch model directory.
+    """
+    from executorch import version as executorch_version
+    from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
+    from executorch.exir import to_edge_transform_and_lower
+
+    LOGGER.info(f"\n{prefix} starting export with ExecuTorch {executorch_version.__version__}...")
+
+    file = Path(file)
+    output_dir = Path(str(file).replace(file.suffix, "_executorch_model"))
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    pte_file = output_dir / file.with_suffix(".pte").name
+    et_program = to_edge_transform_and_lower(
+        torch.export.export(model, (sample_input,)),
+        partitioner=[XnnpackPartitioner()],
+    ).to_executorch()
+    pte_file.write_bytes(et_program.buffer)
+
+    if metadata is not None:
+        YAML.save(output_dir / "metadata.yaml", metadata)
+
+    return str(output_dir)

ultralytics/utils/export/imx.py

@@ -0,0 +1,345 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import types
+from pathlib import Path
+from shutil import which
+
+import numpy as np
+import torch
+
+from ultralytics.nn.modules import Detect, Pose, Segment
+from ultralytics.utils import LOGGER, WINDOWS
+from ultralytics.utils.patches import onnx_export_patch
+from ultralytics.utils.tal import make_anchors
+from ultralytics.utils.torch_utils import copy_attr
+
+# Configuration for Model Compression Toolkit (MCT) quantization
+MCT_CONFIG = {
+    "YOLO11": {
+        "detect": {
+            "layer_names": ["sub", "mul_2", "add_14", "cat_19"],
+            "weights_memory": 2585350.2439,
+            "n_layers": {238, 239},
+        },
+        "pose": {
+            "layer_names": ["sub", "mul_2", "add_14", "cat_21", "cat_22", "mul_4", "add_15"],
+            "weights_memory": 2437771.67,
+            "n_layers": {257, 258},
+        },
+        "classify": {"layer_names": [], "weights_memory": np.inf, "n_layers": {112}},
+        "segment": {
+            "layer_names": ["sub", "mul_2", "add_14", "cat_21"],
+            "weights_memory": 2466604.8,
+            "n_layers": {265, 266},
+        },
+    },
+    "YOLOv8": {
+        "detect": {
+            "layer_names": ["sub", "mul", "add_6", "cat_15"],
+            "weights_memory": 2550540.8,
+            "n_layers": {168, 169},
+        },
+        "pose": {
+            "layer_names": ["add_7", "mul_2", "cat_17", "mul", "sub", "add_6", "cat_18"],
+            "weights_memory": 2482451.85,
+            "n_layers": {187, 188},
+        },
+        "classify": {"layer_names": [], "weights_memory": np.inf, "n_layers": {73}},
+        "segment": {
+            "layer_names": ["sub", "mul", "add_6", "cat_17"],
+            "weights_memory": 2580060.0,
+            "n_layers": {195, 196},
+        },
+    },
+}
+
+
+class FXModel(torch.nn.Module):
+    """A custom model class for torch.fx compatibility.
+
+    This class extends `torch.nn.Module` and is designed to ensure compatibility with torch.fx for tracing and graph
+    manipulation. It copies attributes from an existing model and explicitly sets the model attribute to ensure proper
+    copying.
+
+    Attributes:
+        model (nn.Module): The original model's layers.
+        imgsz (tuple[int, int]): The input image size (height, width).
+    """
+
+    def __init__(self, model, imgsz=(640, 640)):
+        """Initialize the FXModel.
+
+        Args:
+            model (nn.Module): The original model to wrap for torch.fx compatibility.
+            imgsz (tuple[int, int]): The input image size (height, width). Default is (640, 640).
+        """
+        super().__init__()
+        copy_attr(self, model)
+        # Explicitly set `model` since `copy_attr` somehow does not copy it.
+        self.model = model.model
+        self.imgsz = imgsz
+
+    def forward(self, x):
+        """Forward pass through the model.
+
+        This method performs the forward pass through the model, handling the dependencies between layers and saving
+        intermediate outputs.
+
+        Args:
+            x (torch.Tensor): The input tensor to the model.
+
+        Returns:
+            (torch.Tensor): The output tensor from the model.
+        """
+        y = []  # outputs
+        for m in self.model:
+            if m.f != -1:  # if not from previous layer
+                # from earlier layers
+                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]
+            if isinstance(m, Detect):
+                m._inference = types.MethodType(_inference, m)  # bind method to Detect
+                m.anchors, m.strides = (
+                    x.transpose(0, 1)
+                    for x in make_anchors(
+                        torch.cat([s / m.stride.unsqueeze(-1) for s in self.imgsz], dim=1), m.stride, 0.5
+                    )
+                )
+            if type(m) is Pose:
+                m.forward = types.MethodType(pose_forward, m)  # bind method to Pose
+            if type(m) is Segment:
+                m.forward = types.MethodType(segment_forward, m)  # bind method to Segment
+            x = m(x)  # run
+            y.append(x)  # save output
+        return x
+
+
+def _inference(self, x: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+    """Decode boxes and cls scores for imx object detection."""
+    dbox = self.decode_bboxes(self.dfl(x["boxes"]), self.anchors.unsqueeze(0)) * self.strides
+    return dbox.transpose(1, 2), x["scores"].sigmoid().permute(0, 2, 1)
+
+
+def pose_forward(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Forward pass for imx pose estimation, including keypoint decoding."""
+    bs = x[0].shape[0]  # batch size
+    nk_out = getattr(self, "nk_output", self.nk)
+    kpt = torch.cat([self.cv4[i](x[i]).view(bs, nk_out, -1) for i in range(self.nl)], -1)
+
+    # If using Pose26 with 5 dims, convert to 3 dims for export
+    if hasattr(self, "nk_output") and self.nk_output != self.nk:
+        spatial = kpt.shape[-1]
+        kpt = kpt.view(bs, self.kpt_shape[0], self.kpt_shape[1] + 2, spatial)
+        kpt = kpt[:, :, :-2, :]  # Remove sigma_x, sigma_y
+        kpt = kpt.view(bs, self.nk, spatial)
+    x = Detect.forward(self, x)
+    pred_kpt = self.kpts_decode(kpt)
+    return *x, pred_kpt.permute(0, 2, 1)
+
+
+def segment_forward(self, x: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Forward pass for imx segmentation."""
+    p = self.proto(x[0])  # mask protos
+    bs = p.shape[0]  # batch size
+    mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2)  # mask coefficients
+    x = Detect.forward(self, x)
+    return *x, mc.transpose(1, 2), p
+
+
+class NMSWrapper(torch.nn.Module):
+    """Wrap PyTorch Module with multiclass_nms layer from edge-mdt-cl."""
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        score_threshold: float = 0.001,
+        iou_threshold: float = 0.7,
+        max_detections: int = 300,
+        task: str = "detect",
+    ):
+        """Initialize NMSWrapper with PyTorch Module and NMS parameters.
+
+        Args:
+            model (torch.nn.Module): Model instance.
+            score_threshold (float): Score threshold for non-maximum suppression.
+            iou_threshold (float): Intersection over union threshold for non-maximum suppression.
+            max_detections (int): The number of detections to return.
+            task (str): Task type, one of 'detect', 'pose', or 'segment'.
+        """
+        super().__init__()
+        self.model = model
+        self.score_threshold = score_threshold
+        self.iou_threshold = iou_threshold
+        self.max_detections = max_detections
+        self.task = task
+
+    def forward(self, images):
+        """Forward pass with model inference and NMS post-processing."""
+        from edgemdt_cl.pytorch.nms.nms_with_indices import multiclass_nms_with_indices
+
+        # model inference
+        outputs = self.model(images)
+        boxes, scores = outputs[0], outputs[1]
+        nms_outputs = multiclass_nms_with_indices(
+            boxes=boxes,
+            scores=scores,
+            score_threshold=self.score_threshold,
+            iou_threshold=self.iou_threshold,
+            max_detections=self.max_detections,
+        )
+        if self.task == "pose":
+            kpts = outputs[2]  # (bs, max_detections, kpts 17*3)
+            out_kpts = torch.gather(kpts, 1, nms_outputs.indices.unsqueeze(-1).expand(-1, -1, kpts.size(-1)))
+            return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, out_kpts
+        if self.task == "segment":
+            mc, proto = outputs[2], outputs[3]
+            out_mc = torch.gather(mc, 1, nms_outputs.indices.unsqueeze(-1).expand(-1, -1, mc.size(-1)))
+            return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, out_mc, proto
+        return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, nms_outputs.n_valid
+
+
+def torch2imx(
+    model: torch.nn.Module,
+    file: Path | str,
+    conf: float,
+    iou: float,
+    max_det: int,
+    metadata: dict | None = None,
+    gptq: bool = False,
+    dataset=None,
+    prefix: str = "",
+):
+    """Export YOLO model to IMX format for deployment on Sony IMX500 devices.
+
+    This function quantizes a YOLO model using Model Compression Toolkit (MCT) and exports it to IMX format compatible
+    with Sony IMX500 edge devices. It supports both YOLOv8n and YOLO11n models for detection, segmentation, pose
+    estimation, and classification tasks.
+
+    Args:
+        model (torch.nn.Module): The YOLO model to export. Must be YOLOv8n or YOLO11n.
+        file (Path | str): Output file path for the exported model.
+        conf (float): Confidence threshold for NMS post-processing.
+        iou (float): IoU threshold for NMS post-processing.
+        max_det (int): Maximum number of detections to return.
+        metadata (dict | None, optional): Metadata to embed in the ONNX model. Defaults to None.
+        gptq (bool, optional): Whether to use Gradient-Based Post Training Quantization. If False, uses standard Post
+            Training Quantization. Defaults to False.
+        dataset (optional): Representative dataset for quantization calibration. Defaults to None.
+        prefix (str, optional): Logging prefix string. Defaults to "".
+
+    Returns:
+        (Path): Path to the exported IMX model directory.
+
+    Raises:
+        ValueError: If the model is not a supported YOLOv8n or YOLO11n variant.
+
+    Examples:
+        >>> from ultralytics import YOLO
+        >>> model = YOLO("yolo11n.pt")
+        >>> path = torch2imx(model, "model.imx", conf=0.25, iou=0.7, max_det=300)
+
+    Notes:
+        - Requires model_compression_toolkit, onnx, edgemdt_tpc, and edge-mdt-cl packages
+        - Only supports YOLOv8n and YOLO11n models (detection, segmentation, pose, and classification tasks)
+        - Output includes quantized ONNX model, IMX binary, and labels.txt file
+    """
+    import model_compression_toolkit as mct
+    import onnx
+    from edgemdt_tpc import get_target_platform_capabilities
+
+    LOGGER.info(f"\n{prefix} starting export with model_compression_toolkit {mct.__version__}...")
+
+    def representative_dataset_gen(dataloader=dataset):
+        for batch in dataloader:
+            img = batch["img"]
+            img = img / 255.0
+            yield [img]
+
+    # NOTE: need tpc_version to be "4.0" for IMX500 Pose estimation models
+    tpc = get_target_platform_capabilities(tpc_version="4.0", device_type="imx500")
+
+    bit_cfg = mct.core.BitWidthConfig()
+    mct_config = MCT_CONFIG["YOLO11" if "C2PSA" in model.__str__() else "YOLOv8"][model.task]
+
+    # Check if the model has the expected number of layers
+    if len(list(model.modules())) not in mct_config["n_layers"]:
+        raise ValueError("IMX export only supported for YOLOv8n and YOLO11n models.")
+
+    for layer_name in mct_config["layer_names"]:
+        bit_cfg.set_manual_activation_bit_width([mct.core.common.network_editors.NodeNameFilter(layer_name)], 16)
+
+    config = mct.core.CoreConfig(
+        mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=10),
+        quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True),
+        bit_width_config=bit_cfg,
+    )
+
+    resource_utilization = mct.core.ResourceUtilization(weights_memory=mct_config["weights_memory"])
+
+    quant_model = (
+        mct.gptq.pytorch_gradient_post_training_quantization(  # Perform Gradient-Based Post Training Quantization
+            model=model,
+            representative_data_gen=representative_dataset_gen,
+            target_resource_utilization=resource_utilization,
+            gptq_config=mct.gptq.get_pytorch_gptq_config(
+                n_epochs=1000, use_hessian_based_weights=False, use_hessian_sample_attention=False
+            ),
+            core_config=config,
+            target_platform_capabilities=tpc,
+        )[0]
+        if gptq
+        else mct.ptq.pytorch_post_training_quantization(  # Perform post training quantization
+            in_module=model,
+            representative_data_gen=representative_dataset_gen,
+            target_resource_utilization=resource_utilization,
+            core_config=config,
+            target_platform_capabilities=tpc,
+        )[0]
+    )
+
+    if model.task != "classify":
+        quant_model = NMSWrapper(
+            model=quant_model,
+            score_threshold=conf or 0.001,
+            iou_threshold=iou,
+            max_detections=max_det,
+            task=model.task,
+        )
+
+    f = Path(str(file).replace(file.suffix, "_imx_model"))
+    f.mkdir(exist_ok=True)
+    onnx_model = f / Path(str(file.name).replace(file.suffix, "_imx.onnx"))  # js dir
+
+    with onnx_export_patch():
+        mct.exporter.pytorch_export_model(
+            model=quant_model, save_model_path=onnx_model, repr_dataset=representative_dataset_gen
+        )
+
+    model_onnx = onnx.load(onnx_model)  # load onnx model
+    for k, v in metadata.items():
+        meta = model_onnx.metadata_props.add()
+        meta.key, meta.value = k, str(v)
+
+    onnx.save(model_onnx, onnx_model)
+
+    # Find imxconv-pt binary - check venv bin directory first, then PATH
+    bin_dir = Path(sys.executable).parent
+    imxconv = bin_dir / ("imxconv-pt.exe" if WINDOWS else "imxconv-pt")
+    if not imxconv.exists():
+        imxconv = which("imxconv-pt")  # fallback to PATH
+    if not imxconv:
+        raise FileNotFoundError("imxconv-pt not found. Install with: pip install imx500-converter[pt]")
+
+    subprocess.run(
+        [str(imxconv), "-i", str(onnx_model), "-o", str(f), "--no-input-persistency", "--overwrite-output"],
+        check=True,
+    )
+
+    # Needed for imx models.
+    with open(f / "labels.txt", "w", encoding="utf-8") as file:
+        file.writelines([f"{name}\n" for _, name in model.names.items()])
+
+    return f

ultralytics/utils/export/tensorflow.py

@@ -0,0 +1,231 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from functools import partial
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from ultralytics.nn.modules import Detect, Pose, Pose26
+from ultralytics.utils import LOGGER
+from ultralytics.utils.downloads import attempt_download_asset
+from ultralytics.utils.files import spaces_in_path
+from ultralytics.utils.tal import make_anchors
+
+
+def tf_wrapper(model: torch.nn.Module) -> torch.nn.Module:
+    """A wrapper for TensorFlow export compatibility (TF-specific handling is now in head modules)."""
+    for m in model.modules():
+        if not isinstance(m, Detect):
+            continue
+        import types
+
+        m._get_decode_boxes = types.MethodType(_tf_decode_boxes, m)
+        if isinstance(m, Pose):
+            m.kpts_decode = types.MethodType(partial(_tf_kpts_decode, is_pose26=type(m) is Pose26), m)
+    return model
+
+
+def _tf_decode_boxes(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
+    """Decode bounding boxes for TensorFlow export."""
+    shape = x["feats"][0].shape  # BCHW
+    boxes = x["boxes"]
+    if self.format != "imx" and (self.dynamic or self.shape != shape):
+        self.anchors, self.strides = (a.transpose(0, 1) for a in make_anchors(x["feats"], self.stride, 0.5))
+        self.shape = shape
+    grid_h, grid_w = shape[2:4]
+    grid_size = torch.tensor([grid_w, grid_h, grid_w, grid_h], device=boxes.device).reshape(1, 4, 1)
+    norm = self.strides / (self.stride[0] * grid_size)
+    dbox = self.decode_bboxes(self.dfl(boxes) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
+    return dbox
+
+
+def _tf_kpts_decode(self, kpts: torch.Tensor, is_pose26: bool = False) -> torch.Tensor:
+    """Decode keypoints for TensorFlow export."""
+    ndim = self.kpt_shape[1]
+    bs = kpts.shape[0]
+    # Precompute normalization factor to increase numerical stability
+    y = kpts.view(bs, *self.kpt_shape, -1)
+    grid_h, grid_w = self.shape[2:4]
+    grid_size = torch.tensor([grid_w, grid_h], device=y.device).reshape(1, 2, 1)
+    norm = self.strides / (self.stride[0] * grid_size)
+    a = ((y[:, :, :2] + self.anchors) if is_pose26 else (y[:, :, :2] * 2.0 + (self.anchors - 0.5))) * norm
+    if ndim == 3:
+        a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
+    return a.view(bs, self.nk, -1)
+
+
+def onnx2saved_model(
+    onnx_file: str,
+    output_dir: Path,
+    int8: bool = False,
+    images: np.ndarray = None,
+    disable_group_convolution: bool = False,
+    prefix="",
+):
+    """Convert an ONNX model to TensorFlow SavedModel format using onnx2tf.
+
+    Args:
+        onnx_file (str): ONNX file path.
+        output_dir (Path): Output directory path for the SavedModel.
+        int8 (bool, optional): Enable INT8 quantization. Defaults to False.
+        images (np.ndarray, optional): Calibration images for INT8 quantization in BHWC format.
+        disable_group_convolution (bool, optional): Disable group convolution optimization. Defaults to False.
+        prefix (str, optional): Logging prefix. Defaults to "".
+
+    Returns:
+        (keras.Model): Converted Keras model.
+
+    Notes:
+        - Requires onnx2tf package. Downloads calibration data if INT8 quantization is enabled.
+        - Removes temporary files and renames quantized models after conversion.
+    """
+    # Pre-download calibration file to fix https://github.com/PINTO0309/onnx2tf/issues/545
+    onnx2tf_file = Path("calibration_image_sample_data_20x128x128x3_float32.npy")
+    if not onnx2tf_file.exists():
+        attempt_download_asset(f"{onnx2tf_file}.zip", unzip=True, delete=True)
+    np_data = None
+    if int8:
+        tmp_file = output_dir / "tmp_tflite_int8_calibration_images.npy"  # int8 calibration images file
+        if images is not None:
+            output_dir.mkdir(parents=True, exist_ok=True)
+            np.save(str(tmp_file), images)  # BHWC
+            np_data = [["images", tmp_file, [[[[0, 0, 0]]]], [[[[255, 255, 255]]]]]]
+
+    # Patch onnx.helper for onnx_graphsurgeon compatibility with ONNX>=1.17
+    # The float32_to_bfloat16 function was removed in ONNX 1.17, but onnx_graphsurgeon still uses it
+    import onnx.helper
+
+    if not hasattr(onnx.helper, "float32_to_bfloat16"):
+        import struct
+
+        def float32_to_bfloat16(fval):
+            """Convert float32 to bfloat16 (truncates lower 16 bits of mantissa)."""
+            ival = struct.unpack("=I", struct.pack("=f", fval))[0]
+            return ival >> 16
+
+        onnx.helper.float32_to_bfloat16 = float32_to_bfloat16
+
+    import onnx2tf  # scoped for after ONNX export for reduced conflict during import
+
+    LOGGER.info(f"{prefix} starting TFLite export with onnx2tf {onnx2tf.__version__}...")
+    keras_model = onnx2tf.convert(
+        input_onnx_file_path=onnx_file,
+        output_folder_path=str(output_dir),
+        not_use_onnxsim=True,
+        verbosity="error",  # note INT8-FP16 activation bug https://github.com/ultralytics/ultralytics/issues/15873
+        output_integer_quantized_tflite=int8,
+        custom_input_op_name_np_data_path=np_data,
+        enable_batchmatmul_unfold=True and not int8,  # fix lower no. of detected objects on GPU delegate
+        output_signaturedefs=True,  # fix error with Attention block group convolution
+        disable_group_convolution=disable_group_convolution,  # fix error with group convolution
+    )
+
+    # Remove/rename TFLite models
+    if int8:
+        tmp_file.unlink(missing_ok=True)
+        for file in output_dir.rglob("*_dynamic_range_quant.tflite"):
+            file.rename(file.with_name(file.stem.replace("_dynamic_range_quant", "_int8") + file.suffix))
+        for file in output_dir.rglob("*_integer_quant_with_int16_act.tflite"):
+            file.unlink()  # delete extra fp16 activation TFLite files
+    return keras_model
+
+
+def keras2pb(keras_model, file: Path, prefix=""):
+    """Convert a Keras model to TensorFlow GraphDef (.pb) format.
+
+    Args:
+        keras_model (keras.Model): Keras model to convert to frozen graph format.
+        file (Path): Output file path (suffix will be changed to .pb).
+        prefix (str, optional): Logging prefix. Defaults to "".
+
+    Notes:
+        Creates a frozen graph by converting variables to constants for inference optimization.
+    """
+    import tensorflow as tf
+    from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
+
+    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
+    m = tf.function(lambda x: keras_model(x))  # full model
+    m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype))
+    frozen_func = convert_variables_to_constants_v2(m)
+    frozen_func.graph.as_graph_def()
+    tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(file.parent), name=file.name, as_text=False)
+
+
+def tflite2edgetpu(tflite_file: str | Path, output_dir: str | Path, prefix: str = ""):
+    """Convert a TensorFlow Lite model to Edge TPU format using the Edge TPU compiler.
+
+    Args:
+        tflite_file (str | Path): Path to the input TensorFlow Lite (.tflite) model file.
+        output_dir (str | Path): Output directory path for the compiled Edge TPU model.
+        prefix (str, optional): Logging prefix. Defaults to "".
+
+    Notes:
+        Requires the Edge TPU compiler to be installed. The function compiles the TFLite model
+        for optimal performance on Google's Edge TPU hardware accelerator.
+    """
+    import subprocess
+
+    cmd = (
+        "edgetpu_compiler "
+        f'--out_dir "{output_dir}" '
+        "--show_operations "
+        "--search_delegate "
+        "--delegate_search_step 30 "
+        "--timeout_sec 180 "
+        f'"{tflite_file}"'
+    )
+    LOGGER.info(f"{prefix} running '{cmd}'")
+    subprocess.run(cmd, shell=True)
+
+
+def pb2tfjs(pb_file: str, output_dir: str, half: bool = False, int8: bool = False, prefix: str = ""):
+    """Convert a TensorFlow GraphDef (.pb) model to TensorFlow.js format.
+
+    Args:
+        pb_file (str): Path to the input TensorFlow GraphDef (.pb) model file.
+        output_dir (str): Output directory path for the converted TensorFlow.js model.
+        half (bool, optional): Enable FP16 quantization. Defaults to False.
+        int8 (bool, optional): Enable INT8 quantization. Defaults to False.
+        prefix (str, optional): Logging prefix. Defaults to "".
+
+    Notes:
+        Requires tensorflowjs package. Uses tensorflowjs_converter command-line tool for conversion.
+        Handles spaces in file paths and warns if output directory contains spaces.
+    """
+    import subprocess
+
+    import tensorflow as tf
+    import tensorflowjs as tfjs
+
+    LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
+
+    gd = tf.Graph().as_graph_def()  # TF GraphDef
+    with open(pb_file, "rb") as file:
+        gd.ParseFromString(file.read())
+    outputs = ",".join(gd_outputs(gd))
+    LOGGER.info(f"\n{prefix} output node names: {outputs}")
+
+    quantization = "--quantize_float16" if half else "--quantize_uint8" if int8 else ""
+    with spaces_in_path(pb_file) as fpb_, spaces_in_path(output_dir) as f_:  # exporter cannot handle spaces in paths
+        cmd = (
+            "tensorflowjs_converter "
+            f'--input_format=tf_frozen_model {quantization} --output_node_names={outputs} "{fpb_}" "{f_}"'
+        )
+        LOGGER.info(f"{prefix} running '{cmd}'")
+        subprocess.run(cmd, shell=True)
+
+    if " " in output_dir:
+        LOGGER.warning(f"{prefix} your model may not work correctly with spaces in path '{output_dir}'.")
+
+
+def gd_outputs(gd):
+    """Return TensorFlow GraphDef model output node names."""
+    name_list, input_list = [], []
+    for node in gd.node:  # tensorflow.core.framework.node_def_pb2.NodeDef
+        name_list.append(node.name)
+        input_list.extend(node.input)
+    return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))

ultralytics/utils/files.py

@@ -0,0 +1,219 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import contextlib
+import glob
+import os
+import shutil
+import tempfile
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+
+
+class WorkingDirectory(contextlib.ContextDecorator):
+    """A context manager and decorator for temporarily changing the working directory.
+
+    This class allows for the temporary change of the working directory using a context manager or decorator. It ensures
+    that the original working directory is restored after the context or decorated function completes.
+
+    Attributes:
+        dir (Path | str): The new directory to switch to.
+        cwd (Path): The original current working directory before the switch.
+
+    Methods:
+        __enter__: Changes the current directory to the specified directory.
+        __exit__: Restores the original working directory on context exit.
+
+    Examples:
+        Using as a context manager:
+        >>> with WorkingDirectory("/path/to/new/dir"):
+        ...     # Perform operations in the new directory
+        ...     pass
+
+        Using as a decorator:
+        >>> @WorkingDirectory("/path/to/new/dir")
+        ... def some_function():
+        ...     # Perform operations in the new directory
+        ...     pass
+    """
+
+    def __init__(self, new_dir: str | Path):
+        """Initialize the WorkingDirectory context manager with the target directory."""
+        self.dir = new_dir  # new dir
+        self.cwd = Path.cwd().resolve()  # current dir
+
+    def __enter__(self):
+        """Change the current working directory to the specified directory upon entering the context."""
+        os.chdir(self.dir)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Restore the original working directory when exiting the context."""
+        os.chdir(self.cwd)
+
+
+@contextmanager
+def spaces_in_path(path: str | Path):
+    """Context manager to handle paths with spaces in their names.
+
+    If a path contains spaces, it replaces them with underscores, copies the file/directory to the new path, executes
+    the context code block, then copies the file/directory back to its original location.
+
+    Args:
+        path (str | Path): The original path that may contain spaces.
+
+    Yields:
+        (Path | str): Temporary path with any spaces replaced by underscores.
+
+    Examples:
+        >>> with spaces_in_path("/path/with spaces") as new_path:
+        ...     # Your code here
+        ...     pass
+    """
+    # If path has spaces, replace them with underscores
+    if " " in str(path):
+        string = isinstance(path, str)  # input type
+        path = Path(path)
+
+        # Create a temporary directory and construct the new path
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_path = Path(tmp_dir) / path.name.replace(" ", "_")
+
+            # Copy file/directory
+            if path.is_dir():
+                shutil.copytree(path, tmp_path)
+            elif path.is_file():
+                tmp_path.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copy2(path, tmp_path)
+
+            try:
+                # Yield the temporary path
+                yield str(tmp_path) if string else tmp_path
+
+            finally:
+                # Copy file/directory back
+                if tmp_path.is_dir():
+                    shutil.copytree(tmp_path, path, dirs_exist_ok=True)
+                elif tmp_path.is_file():
+                    shutil.copy2(tmp_path, path)  # Copy back the file
+
+    else:
+        # If there are no spaces, just yield the original path
+        yield path
+
+
+def increment_path(path: str | Path, exist_ok: bool = False, sep: str = "", mkdir: bool = False) -> Path:
+    """Increment a file or directory path, i.e., runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
+
+    If the path exists and `exist_ok` is not True, the path will be incremented by appending a number and `sep` to the
+    end of the path. If the path is a file, the file extension will be preserved. If the path is a directory, the number
+    will be appended directly to the end of the path.
+
+    Args:
+        path (str | Path): Path to increment.
+        exist_ok (bool, optional): If True, the path will not be incremented and returned as-is.
+        sep (str, optional): Separator to use between the path and the incrementation number.
+        mkdir (bool, optional): Create a directory if it does not exist.
+
+    Returns:
+        (Path): Incremented path.
+
+    Examples:
+        Increment a directory path:
+        >>> from pathlib import Path
+        >>> path = Path("runs/exp")
+        >>> new_path = increment_path(path)
+        >>> print(new_path)
+        runs/exp2
+
+        Increment a file path:
+        >>> path = Path("runs/exp/results.txt")
+        >>> new_path = increment_path(path)
+        >>> print(new_path)
+        runs/exp/results2.txt
+    """
+    path = Path(path)  # os-agnostic
+    if path.exists() and not exist_ok:
+        path, suffix = (path.with_suffix(""), path.suffix) if path.is_file() else (path, "")
+
+        # Method 1
+        for n in range(2, 9999):
+            p = f"{path}{sep}{n}{suffix}"  # increment path
+            if not os.path.exists(p):
+                break
+        path = Path(p)
+
+    if mkdir:
+        path.mkdir(parents=True, exist_ok=True)  # make directory
+
+    return path
+
+
+def file_age(path: str | Path = __file__) -> int:
+    """Return days since the last modification of the specified file."""
+    dt = datetime.now() - datetime.fromtimestamp(Path(path).stat().st_mtime)  # delta
+    return dt.days  # + dt.seconds / 86400  # fractional days
+
+
+def file_date(path: str | Path = __file__) -> str:
+    """Return the file modification date in 'YYYY-M-D' format."""
+    t = datetime.fromtimestamp(Path(path).stat().st_mtime)
+    return f"{t.year}-{t.month}-{t.day}"
+
+
+def file_size(path: str | Path) -> float:
+    """Return the size of a file or directory in mebibytes (MiB)."""
+    if isinstance(path, (str, Path)):
+        mb = 1 << 20  # bytes to MiB (1024 ** 2)
+        path = Path(path)
+        if path.is_file():
+            return path.stat().st_size / mb
+        elif path.is_dir():
+            return sum(f.stat().st_size for f in path.glob("**/*") if f.is_file()) / mb
+    return 0.0
+
+
+def get_latest_run(search_dir: str = ".") -> str:
+    """Return the path to the most recent 'last.pt' file in the specified directory for resuming training."""
+    last_list = glob.glob(f"{search_dir}/**/last*.pt", recursive=True)
+    return max(last_list, key=os.path.getctime) if last_list else ""
+
+
+def update_models(model_names: tuple = ("yolo26n.pt",), source_dir: Path = Path("."), update_names: bool = False):
+    """Update and re-save specified YOLO models in an 'updated_models' subdirectory.
+
+    Args:
+        model_names (tuple, optional): Model filenames to update.
+        source_dir (Path, optional): Directory containing models and target subdirectory.
+        update_names (bool, optional): Update model names from a data YAML.
+
+    Examples:
+        Update specified YOLO models and save them in 'updated_models' subdirectory:
+        >>> from ultralytics.utils.files import update_models
+        >>> model_names = ("yolo26n.pt", "yolo11s.pt")
+        >>> update_models(model_names, source_dir=Path("/models"), update_names=True)
+    """
+    from ultralytics import YOLO
+    from ultralytics.nn.autobackend import default_class_names
+    from ultralytics.utils import LOGGER
+
+    target_dir = source_dir / "updated_models"
+    target_dir.mkdir(parents=True, exist_ok=True)  # Ensure target directory exists
+
+    for model_name in model_names:
+        model_path = source_dir / model_name
+        LOGGER.info(f"Loading model from {model_path}")
+
+        # Load model
+        model = YOLO(model_path)
+        model.half()
+        if update_names:  # update model names from a dataset YAML
+            model.model.names = default_class_names("coco8.yaml")
+
+        # Define new save path
+        save_path = target_dir / model_name
+
+        # Save model using model.save()
+        LOGGER.info(f"Re-saving {model_name} model to {save_path}")
+        model.save(save_path)

ultralytics/utils/git.py

@@ -0,0 +1,137 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from functools import cached_property
+from pathlib import Path
+
+
+class GitRepo:
+    """Represent a local Git repository and expose branch, commit, and remote metadata.
+
+    This class discovers the repository root by searching for a .git entry from the given path upward, resolves the
+    actual .git directory (including worktrees), and reads Git metadata directly from on-disk files. It does not invoke
+    the git binary and therefore works in restricted environments. All metadata properties are resolved lazily and
+    cached; construct a new instance to refresh state.
+
+    Attributes:
+        root (Path | None): Repository root directory containing the .git entry; None if not in a repository.
+        gitdir (Path | None): Resolved .git directory path; handles worktrees; None if unresolved.
+        head (str | None): Raw contents of HEAD; a SHA for detached HEAD or "ref: <refname>" for branch heads.
+        is_repo (bool): Whether the provided path resides inside a Git repository.
+        branch (str | None): Current branch name when HEAD points to a branch; None for detached HEAD or non-repo.
+        commit (str | None): Current commit SHA for HEAD; None if not determinable.
+        origin (str | None): URL of the "origin" remote as read from gitdir/config; None if unset or unavailable.
+
+    Examples:
+        Initialize from the current working directory and read metadata
+        >>> from pathlib import Path
+        >>> repo = GitRepo(Path.cwd())
+        >>> repo.is_repo
+        True
+        >>> repo.branch, repo.commit[:7], repo.origin
+        ('main', '1a2b3c4', 'https://example.com/owner/repo.git')
+
+    Notes:
+        - Resolves metadata by reading files: HEAD, packed-refs, and config; no subprocess calls are used.
+        - Caches properties on first access using cached_property; recreate the object to reflect repository changes.
+    """
+
+    def __init__(self, path: Path = Path(__file__).resolve()):
+        """Initialize a Git repository context by discovering the repository root from a starting path.
+
+        Args:
+            path (Path, optional): File or directory path used as the starting point to locate the repository root.
+        """
+        self.root = self._find_root(path)
+        self.gitdir = self._gitdir(self.root) if self.root else None
+
+    @staticmethod
+    def _find_root(p: Path) -> Path | None:
+        """Return repo root or None."""
+        return next((d for d in [p, *list(p.parents)] if (d / ".git").exists()), None)
+
+    @staticmethod
+    def _gitdir(root: Path) -> Path | None:
+        """Resolve actual .git directory (handles worktrees)."""
+        g = root / ".git"
+        if g.is_dir():
+            return g
+        if g.is_file():
+            t = g.read_text(errors="ignore").strip()
+            if t.startswith("gitdir:"):
+                return (root / t.split(":", 1)[1].strip()).resolve()
+        return None
+
+    @staticmethod
+    def _read(p: Path | None) -> str | None:
+        """Read and strip file if exists."""
+        return p.read_text(errors="ignore").strip() if p and p.exists() else None
+
+    @cached_property
+    def head(self) -> str | None:
+        """HEAD file contents."""
+        return self._read(self.gitdir / "HEAD" if self.gitdir else None)
+
+    def _ref_commit(self, ref: str) -> str | None:
+        """Commit for ref (handles packed-refs)."""
+        rf = self.gitdir / ref
+        if s := self._read(rf):
+            return s
+        pf = self.gitdir / "packed-refs"
+        b = pf.read_bytes().splitlines() if pf.exists() else []
+        tgt = ref.encode()
+        for line in b:
+            if line[:1] in (b"#", b"^") or b" " not in line:
+                continue
+            sha, name = line.split(b" ", 1)
+            if name.strip() == tgt:
+                return sha.decode()
+        return None
+
+    @property
+    def is_repo(self) -> bool:
+        """True if inside a git repo."""
+        return self.gitdir is not None
+
+    @cached_property
+    def branch(self) -> str | None:
+        """Current branch or None."""
+        if not self.is_repo or not self.head or not self.head.startswith("ref: "):
+            return None
+        ref = self.head[5:].strip()
+        return ref[len("refs/heads/") :] if ref.startswith("refs/heads/") else ref
+
+    @cached_property
+    def commit(self) -> str | None:
+        """Current commit SHA or None."""
+        if not self.is_repo or not self.head:
+            return None
+        return self._ref_commit(self.head[5:].strip()) if self.head.startswith("ref: ") else self.head
+
+    @cached_property
+    def origin(self) -> str | None:
+        """Origin URL or None."""
+        if not self.is_repo:
+            return None
+        cfg = self.gitdir / "config"
+        remote, url = None, None
+        for s in (self._read(cfg) or "").splitlines():
+            t = s.strip()
+            if t.startswith("[") and t.endswith("]"):
+                remote = t.lower()
+            elif t.lower().startswith("url =") and remote == '[remote "origin"]':
+                url = t.split("=", 1)[1].strip()
+                break
+        return url
+
+
+if __name__ == "__main__":
+    import time
+
+    g = GitRepo()
+    if g.is_repo:
+        t0 = time.perf_counter()
+        print(f"repo={g.root}\nbranch={g.branch}\ncommit={g.commit}\norigin={g.origin}")
+        dt = (time.perf_counter() - t0) * 1000
+        print(f"\n⏱️ Profiling: total {dt:.3f} ms")

ultralytics/utils/instance.py

@@ -0,0 +1,536 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from collections import abc
+from itertools import repeat
+from numbers import Number
+
+import numpy as np
+
+from .ops import ltwh2xywh, ltwh2xyxy, resample_segments, xywh2ltwh, xywh2xyxy, xyxy2ltwh, xyxy2xywh
+
+
+def _ntuple(n):
+    """Create a function that converts input to n-tuple by repeating singleton values."""
+
+    def parse(x):
+        """Parse input to return n-tuple by repeating singleton values n times."""
+        return x if isinstance(x, abc.Iterable) else tuple(repeat(x, n))
+
+    return parse
+
+
+to_2tuple = _ntuple(2)
+to_4tuple = _ntuple(4)
+
+# `xyxy` means left top and right bottom
+# `xywh` means center x, center y and width, height(YOLO format)
+# `ltwh` means left top and width, height(COCO format)
+_formats = ["xyxy", "xywh", "ltwh"]
+
+__all__ = ("Bboxes", "Instances")  # tuple or list
+
+
+class Bboxes:
+    """A class for handling bounding boxes in multiple formats.
+
+    The class supports various bounding box formats like 'xyxy', 'xywh', and 'ltwh' and provides methods for format
+    conversion, scaling, and area calculation. Bounding box data should be provided as numpy arrays.
+
+    Attributes:
+        bboxes (np.ndarray): The bounding boxes stored in a 2D numpy array with shape (N, 4).
+        format (str): The format of the bounding boxes ('xyxy', 'xywh', or 'ltwh').
+
+    Methods:
+        convert: Convert bounding box format from one type to another.
+        areas: Calculate the area of bounding boxes.
+        mul: Multiply bounding box coordinates by scale factor(s).
+        add: Add offset to bounding box coordinates.
+        concatenate: Concatenate multiple Bboxes objects.
+
+    Examples:
+        Create bounding boxes in YOLO format
+        >>> bboxes = Bboxes(np.array([[100, 50, 150, 100]]), format="xywh")
+        >>> bboxes.convert("xyxy")
+        >>> print(bboxes.areas())
+
+    Notes:
+        This class does not handle normalization or denormalization of bounding boxes.
+    """
+
+    def __init__(self, bboxes: np.ndarray, format: str = "xyxy") -> None:
+        """Initialize the Bboxes class with bounding box data in a specified format.
+
+        Args:
+            bboxes (np.ndarray): Array of bounding boxes with shape (N, 4) or (4,).
+            format (str): Format of the bounding boxes, one of 'xyxy', 'xywh', or 'ltwh'.
+        """
+        assert format in _formats, f"Invalid bounding box format: {format}, format must be one of {_formats}"
+        bboxes = bboxes[None, :] if bboxes.ndim == 1 else bboxes
+        assert bboxes.ndim == 2
+        assert bboxes.shape[1] == 4
+        self.bboxes = bboxes
+        self.format = format
+
+    def convert(self, format: str) -> None:
+        """Convert bounding box format from one type to another.
+
+        Args:
+            format (str): Target format for conversion, one of 'xyxy', 'xywh', or 'ltwh'.
+        """
+        assert format in _formats, f"Invalid bounding box format: {format}, format must be one of {_formats}"
+        if self.format == format:
+            return
+        elif self.format == "xyxy":
+            func = xyxy2xywh if format == "xywh" else xyxy2ltwh
+        elif self.format == "xywh":
+            func = xywh2xyxy if format == "xyxy" else xywh2ltwh
+        else:
+            func = ltwh2xyxy if format == "xyxy" else ltwh2xywh
+        self.bboxes = func(self.bboxes)
+        self.format = format
+
+    def areas(self) -> np.ndarray:
+        """Calculate the area of bounding boxes."""
+        return (
+            (self.bboxes[:, 2] - self.bboxes[:, 0]) * (self.bboxes[:, 3] - self.bboxes[:, 1])  # format xyxy
+            if self.format == "xyxy"
+            else self.bboxes[:, 3] * self.bboxes[:, 2]  # format xywh or ltwh
+        )
+
+    def mul(self, scale: int | tuple | list) -> None:
+        """Multiply bounding box coordinates by scale factor(s).
+
+        Args:
+            scale (int | tuple | list): Scale factor(s) for four coordinates. If int, the same scale is applied to all
+                coordinates.
+        """
+        if isinstance(scale, Number):
+            scale = to_4tuple(scale)
+        assert isinstance(scale, (tuple, list))
+        assert len(scale) == 4
+        self.bboxes[:, 0] *= scale[0]
+        self.bboxes[:, 1] *= scale[1]
+        self.bboxes[:, 2] *= scale[2]
+        self.bboxes[:, 3] *= scale[3]
+
+    def add(self, offset: int | tuple | list) -> None:
+        """Add offset to bounding box coordinates.
+
+        Args:
+            offset (int | tuple | list): Offset(s) for four coordinates. If int, the same offset is applied to all
+                coordinates.
+        """
+        if isinstance(offset, Number):
+            offset = to_4tuple(offset)
+        assert isinstance(offset, (tuple, list))
+        assert len(offset) == 4
+        self.bboxes[:, 0] += offset[0]
+        self.bboxes[:, 1] += offset[1]
+        self.bboxes[:, 2] += offset[2]
+        self.bboxes[:, 3] += offset[3]
+
+    def __len__(self) -> int:
+        """Return the number of bounding boxes."""
+        return len(self.bboxes)
+
+    @classmethod
+    def concatenate(cls, boxes_list: list[Bboxes], axis: int = 0) -> Bboxes:
+        """Concatenate a list of Bboxes objects into a single Bboxes object.
+
+        Args:
+            boxes_list (list[Bboxes]): A list of Bboxes objects to concatenate.
+            axis (int, optional): The axis along which to concatenate the bounding boxes.
+
+        Returns:
+            (Bboxes): A new Bboxes object containing the concatenated bounding boxes.
+
+        Notes:
+            The input should be a list or tuple of Bboxes objects.
+        """
+        assert isinstance(boxes_list, (list, tuple))
+        if not boxes_list:
+            return cls(np.empty(0))
+        assert all(isinstance(box, Bboxes) for box in boxes_list)
+
+        if len(boxes_list) == 1:
+            return boxes_list[0]
+        return cls(np.concatenate([b.bboxes for b in boxes_list], axis=axis))
+
+    def __getitem__(self, index: int | np.ndarray | slice) -> Bboxes:
+        """Retrieve a specific bounding box or a set of bounding boxes using indexing.
+
+        Args:
+            index (int | slice | np.ndarray): The index, slice, or boolean array to select the desired bounding boxes.
+
+        Returns:
+            (Bboxes): A new Bboxes object containing the selected bounding boxes.
+
+        Notes:
+            When using boolean indexing, make sure to provide a boolean array with the same length as the number of
+            bounding boxes.
+        """
+        if isinstance(index, int):
+            return Bboxes(self.bboxes[index].reshape(1, -1))
+        b = self.bboxes[index]
+        assert b.ndim == 2, f"Indexing on Bboxes with {index} failed to return a matrix!"
+        return Bboxes(b)
+
+
+class Instances:
+    """Container for bounding boxes, segments, and keypoints of detected objects in an image.
+
+    This class provides a unified interface for handling different types of object annotations including bounding boxes,
+    segmentation masks, and keypoints. It supports various operations like scaling, normalization, clipping, and format
+    conversion.
+
+    Attributes:
+        _bboxes (Bboxes): Internal object for handling bounding box operations.
+        keypoints (np.ndarray): Keypoints with shape (N, 17, 3) in format (x, y, visible).
+        normalized (bool): Flag indicating whether the bounding box coordinates are normalized.
+        segments (np.ndarray): Segments array with shape (N, M, 2) after resampling.
+
+    Methods:
+        convert_bbox: Convert bounding box format.
+        scale: Scale coordinates by given factors.
+        denormalize: Convert normalized coordinates to absolute coordinates.
+        normalize: Convert absolute coordinates to normalized coordinates.
+        add_padding: Add padding to coordinates.
+        flipud: Flip coordinates vertically.
+        fliplr: Flip coordinates horizontally.
+        clip: Clip coordinates to stay within image boundaries.
+        remove_zero_area_boxes: Remove boxes with zero area.
+        update: Update instance variables.
+        concatenate: Concatenate multiple Instances objects.
+
+    Examples:
+        Create instances with bounding boxes and segments
+        >>> instances = Instances(
+        ...     bboxes=np.array([[10, 10, 30, 30], [20, 20, 40, 40]]),
+        ...     segments=[np.array([[5, 5], [10, 10]]), np.array([[15, 15], [20, 20]])],
+        ...     keypoints=np.array([[[5, 5, 1], [10, 10, 1]], [[15, 15, 1], [20, 20, 1]]]),
+        ... )
+    """
+
+    def __init__(
+        self,
+        bboxes: np.ndarray,
+        segments: np.ndarray = None,
+        keypoints: np.ndarray = None,
+        bbox_format: str = "xywh",
+        normalized: bool = True,
+        difficulties: np.ndarray = None,
+        difficulty_levels: np.ndarray = None,
+    ) -> None:
+        """Initialize the Instances object with bounding boxes, segments, and keypoints.
+
+        Args:
+            bboxes (np.ndarray): Bounding boxes with shape (N, 4).
+            segments (np.ndarray, optional): Segmentation masks.
+            keypoints (np.ndarray, optional): Keypoints with shape (N, 17, 3) in format (x, y, visible).
+            bbox_format (str): Format of bboxes.
+            normalized (bool): Whether the coordinates are normalized.
+            difficulties (np.ndarray, optional): Difficulty levels with shape (N,) for ground detection.
+        """
+        self._bboxes = Bboxes(bboxes=bboxes, format=bbox_format)
+        self.keypoints = keypoints
+        self.normalized = normalized
+        self.segments = segments
+        self.difficulties = difficulties
+        self.difficulty_levels = difficulty_levels
+
+    def convert_bbox(self, format: str) -> None:
+        """Convert bounding box format.
+
+        Args:
+            format (str): Target format for conversion, one of 'xyxy', 'xywh', or 'ltwh'.
+        """
+        self._bboxes.convert(format=format)
+
+    @property
+    def bbox_areas(self) -> np.ndarray:
+        """Calculate the area of bounding boxes."""
+        return self._bboxes.areas()
+
+    def scale(self, scale_w: float, scale_h: float, bbox_only: bool = False):
+        """Scale coordinates by given factors.
+
+        Args:
+            scale_w (float): Scale factor for width.
+            scale_h (float): Scale factor for height.
+            bbox_only (bool, optional): Whether to scale only bounding boxes.
+        """
+        self._bboxes.mul(scale=(scale_w, scale_h, scale_w, scale_h))
+        if bbox_only:
+            return
+        self.segments[..., 0] *= scale_w
+        self.segments[..., 1] *= scale_h
+        if self.keypoints is not None:
+            self.keypoints[..., 0] *= scale_w
+            self.keypoints[..., 1] *= scale_h
+
+    def denormalize(self, w: int, h: int) -> None:
+        """Convert normalized coordinates to absolute coordinates.
+
+        Args:
+            w (int): Image width.
+            h (int): Image height.
+        """
+        if not self.normalized:
+            return
+        self._bboxes.mul(scale=(w, h, w, h))
+        self.segments[..., 0] *= w
+        self.segments[..., 1] *= h
+        if self.keypoints is not None:
+            self.keypoints[..., 0] *= w
+            self.keypoints[..., 1] *= h
+        self.normalized = False
+
+    def normalize(self, w: int, h: int) -> None:
+        """Convert absolute coordinates to normalized coordinates.
+
+        Args:
+            w (int): Image width.
+            h (int): Image height.
+        """
+        if self.normalized:
+            return
+        self._bboxes.mul(scale=(1 / w, 1 / h, 1 / w, 1 / h))
+        self.segments[..., 0] /= w
+        self.segments[..., 1] /= h
+        if self.keypoints is not None:
+            self.keypoints[..., 0] /= w
+            self.keypoints[..., 1] /= h
+        self.normalized = True
+
+    def add_padding(self, padw: int, padh: int) -> None:
+        """Add padding to coordinates.
+
+        Args:
+            padw (int): Padding width.
+            padh (int): Padding height.
+        """
+        assert not self.normalized, "you should add padding with absolute coordinates."
+        self._bboxes.add(offset=(padw, padh, padw, padh))
+        self.segments[..., 0] += padw
+        self.segments[..., 1] += padh
+        if self.keypoints is not None:
+            self.keypoints[..., 0] += padw
+            self.keypoints[..., 1] += padh
+
+    def __getitem__(self, index: int | np.ndarray | slice) -> Instances:
+        """Retrieve a specific instance or a set of instances using indexing.
+
+        Args:
+            index (int | slice | np.ndarray): The index, slice, or boolean array to select the desired instances.
+
+        Returns:
+            (Instances): A new Instances object containing the selected boxes, segments, and keypoints if present.
+
+        Notes:
+            When using boolean indexing, make sure to provide a boolean array with the same length as the number of
+            instances.
+        """
+        segments = self.segments[index] if len(self.segments) else self.segments
+        keypoints = self.keypoints[index] if self.keypoints is not None else None
+        difficulties = self.difficulties[index] if self.difficulties is not None else None
+        difficulty_levels = self.difficulty_levels[index] if self.difficulty_levels is not None else None
+        bboxes = self.bboxes[index]
+        bbox_format = self._bboxes.format
+        return Instances(
+            bboxes=bboxes,
+            segments=segments,
+            keypoints=keypoints,
+            bbox_format=bbox_format,
+            normalized=self.normalized,
+            difficulties=difficulties,
+            difficulty_levels=difficulty_levels,
+        )
+
+    def flipud(self, h: int) -> None:
+        """Flip coordinates vertically.
+
+        Args:
+            h (int): Image height.
+        """
+        if self._bboxes.format == "xyxy":
+            y1 = self.bboxes[:, 1].copy()
+            y2 = self.bboxes[:, 3].copy()
+            self.bboxes[:, 1] = h - y2
+            self.bboxes[:, 3] = h - y1
+        else:
+            self.bboxes[:, 1] = h - self.bboxes[:, 1]
+        self.segments[..., 1] = h - self.segments[..., 1]
+        if self.keypoints is not None:
+            self.keypoints[..., 1] = h - self.keypoints[..., 1]
+
+    def fliplr(self, w: int) -> None:
+        """Flip coordinates horizontally.
+
+        Args:
+            w (int): Image width.
+        """
+        if self._bboxes.format == "xyxy":
+            x1 = self.bboxes[:, 0].copy()
+            x2 = self.bboxes[:, 2].copy()
+            self.bboxes[:, 0] = w - x2
+            self.bboxes[:, 2] = w - x1
+        else:
+            self.bboxes[:, 0] = w - self.bboxes[:, 0]
+        self.segments[..., 0] = w - self.segments[..., 0]
+        if self.keypoints is not None:
+            self.keypoints[..., 0] = w - self.keypoints[..., 0]
+
+    def clip(self, w: int, h: int) -> None:
+        """Clip coordinates to stay within image boundaries.
+
+        Args:
+            w (int): Image width.
+            h (int): Image height.
+        """
+        ori_format = self._bboxes.format
+        self.convert_bbox(format="xyxy")
+        self.bboxes[:, [0, 2]] = self.bboxes[:, [0, 2]].clip(0, w)
+        self.bboxes[:, [1, 3]] = self.bboxes[:, [1, 3]].clip(0, h)
+        if ori_format != "xyxy":
+            self.convert_bbox(format=ori_format)
+        self.segments[..., 0] = self.segments[..., 0].clip(0, w)
+        self.segments[..., 1] = self.segments[..., 1].clip(0, h)
+        if self.keypoints is not None:
+            # Set out of bounds visibility to zero
+            self.keypoints[..., 2][
+                (self.keypoints[..., 0] < 0)
+                | (self.keypoints[..., 0] > w)
+                | (self.keypoints[..., 1] < 0)
+                | (self.keypoints[..., 1] > h)
+            ] = 0.0
+            self.keypoints[..., 0] = self.keypoints[..., 0].clip(0, w)
+            self.keypoints[..., 1] = self.keypoints[..., 1].clip(0, h)
+
+    def remove_zero_area_boxes(self) -> np.ndarray:
+        """Remove zero-area boxes, i.e. after clipping some boxes may have zero width or height.
+
+        Returns:
+            (np.ndarray): Boolean array indicating which boxes were kept.
+        """
+        good = self.bbox_areas > 0
+        if not all(good):
+            self._bboxes = self._bboxes[good]
+            if self.segments is not None and len(self.segments):
+                self.segments = self.segments[good]
+            if self.keypoints is not None:
+                self.keypoints = self.keypoints[good]
+            if self.difficulties is not None:
+                self.difficulties = self.difficulties[good]
+            if self.difficulty_levels is not None:
+                self.difficulty_levels = self.difficulty_levels[good]
+        return good
+
+    def update(
+        self,
+        bboxes: np.ndarray,
+        segments: np.ndarray = None,
+        keypoints: np.ndarray = None,
+        difficulties: np.ndarray = None,
+        difficulty_levels: np.ndarray = None,
+    ):
+        """Update instance variables.
+
+        Args:
+            bboxes (np.ndarray): New bounding boxes.
+            segments (np.ndarray, optional): New segments.
+            keypoints (np.ndarray, optional): New keypoints.
+            difficulties (np.ndarray, optional): New difficulty levels.
+        """
+        self._bboxes = Bboxes(bboxes, format=self._bboxes.format)
+        if segments is not None:
+            self.segments = segments
+        if keypoints is not None:
+            self.keypoints = keypoints
+        if difficulties is not None:
+            self.difficulties = difficulties
+        if difficulty_levels is not None:
+            self.difficulty_levels = difficulty_levels
+
+    def __len__(self) -> int:
+        """Return the number of instances."""
+        return len(self.bboxes)
+
+    @classmethod
+    def concatenate(cls, instances_list: list[Instances], axis=0) -> Instances:
+        """Concatenate a list of Instances objects into a single Instances object.
+
+        Args:
+            instances_list (list[Instances]): A list of Instances objects to concatenate.
+            axis (int, optional): The axis along which the arrays will be concatenated.
+
+        Returns:
+            (Instances): A new Instances object containing the concatenated bounding boxes, segments, and keypoints if
+                present.
+
+        Notes:
+            The `Instances` objects in the list should have the same properties, such as the format of the bounding
+            boxes, whether keypoints are present, and if the coordinates are normalized.
+        """
+        assert isinstance(instances_list, (list, tuple))
+        if not instances_list:
+            return cls(np.empty(0))
+        assert all(isinstance(instance, Instances) for instance in instances_list)
+
+        if len(instances_list) == 1:
+            return instances_list[0]
+
+        use_keypoint = instances_list[0].keypoints is not None
+        use_difficulties = instances_list[0].difficulties is not None
+        use_difficulty_levels = instances_list[0].difficulty_levels is not None
+        bbox_format = instances_list[0]._bboxes.format
+        normalized = instances_list[0].normalized
+
+        cat_boxes = np.concatenate([ins.bboxes for ins in instances_list], axis=axis)
+        seg_len = [b.segments.shape[1] for b in instances_list]
+        if len(frozenset(seg_len)) > 1:  # resample segments if there's different length
+            max_len = max(seg_len)
+            cat_segments = np.concatenate(
+                [
+                    resample_segments(list(b.segments), max_len)
+                    if len(b.segments)
+                    else np.zeros((0, max_len, 2), dtype=np.float32)  # re-generating empty segments
+                    for b in instances_list
+                ],
+                axis=axis,
+            )
+        else:
+            cat_segments = np.concatenate([b.segments for b in instances_list], axis=axis)
+        cat_keypoints = np.concatenate([b.keypoints for b in instances_list], axis=axis) if use_keypoint else None
+        cat_difficulties = np.concatenate([b.difficulties for b in instances_list], axis=axis) if use_difficulties else None
+        cat_difficulty_levels = (
+            np.concatenate([b.difficulty_levels for b in instances_list], axis=axis) if use_difficulty_levels else None
+        )
+        return cls(
+            cat_boxes,
+            cat_segments,
+            cat_keypoints,
+            bbox_format,
+            normalized,
+            cat_difficulties,
+            cat_difficulty_levels,
+        )
+
+    @property
+    def bboxes(self) -> np.ndarray:
+        """Return bounding boxes."""
+        return self._bboxes.bboxes
+
+    def __repr__(self) -> str:
+        """Return a string representation of the Instances object."""
+        # Map private to public names and include direct attributes
+        attr_map = {"_bboxes": "bboxes"}
+        parts = []
+        for key, value in self.__dict__.items():
+            name = attr_map.get(key, key)
+            if name == "bboxes":
+                value = self.bboxes  # Use the property
+            if value is not None:
+                parts.append(f"{name}={value!r}")
+        return "Instances({})".format("\n".join(parts))

ultralytics/utils/logger.py

@@ -0,0 +1,506 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import logging
+import shutil
+import sys
+import threading
+import time
+from datetime import datetime
+from pathlib import Path
+
+from ultralytics.utils import LOGGER, MACOS, RANK
+from ultralytics.utils.checks import check_requirements
+
+
+class ConsoleLogger:
+    """Console output capture with batched streaming to file, API, or custom callback.
+
+    Captures stdout/stderr output and streams it with intelligent deduplication and configurable batching.
+
+    Attributes:
+        destination (str | Path | None): Target destination for streaming (URL, Path, or None for callback-only).
+        batch_size (int): Number of lines to batch before flushing (default: 1 for immediate).
+        flush_interval (float): Seconds between automatic flushes (default: 5.0).
+        on_flush (callable | None): Optional callback function called with batched content on flush.
+        active (bool): Whether console capture is currently active.
+
+    Examples:
+        File logging (immediate):
+        >>> logger = ConsoleLogger("training.log")
+        >>> logger.start_capture()
+        >>> print("This will be logged")
+        >>> logger.stop_capture()
+
+        API streaming with batching:
+        >>> logger = ConsoleLogger("https://api.example.com/logs", batch_size=10)
+        >>> logger.start_capture()
+
+        Custom callback with batching:
+        >>> def my_handler(content, line_count, chunk_id):
+        ...     print(f"Received {line_count} lines")
+        >>> logger = ConsoleLogger(on_flush=my_handler, batch_size=5)
+        >>> logger.start_capture()
+    """
+
+    def __init__(self, destination=None, batch_size=1, flush_interval=5.0, on_flush=None):
+        """Initialize console logger with optional batching.
+
+        Args:
+            destination (str | Path | None): API endpoint URL (http/https), local file path, or None.
+            batch_size (int): Lines to accumulate before flush (1 = immediate, higher = batched).
+            flush_interval (float): Max seconds between flushes when batching.
+            on_flush (callable | None): Callback(content: str, line_count: int, chunk_id: int) for custom handling.
+        """
+        self.destination = destination
+        self.is_api = isinstance(destination, str) and destination.startswith(("http://", "https://"))
+        if destination is not None and not self.is_api:
+            self.destination = Path(destination)
+
+        # Batching configuration
+        self.batch_size = max(1, batch_size)
+        self.flush_interval = flush_interval
+        self.on_flush = on_flush
+
+        # Console capture state
+        self.original_stdout = sys.stdout
+        self.original_stderr = sys.stderr
+        self.active = False
+        self._log_handler = None  # Track handler for cleanup
+
+        # Buffer for batching
+        self.buffer = []
+        self.buffer_lock = threading.Lock()
+        self.flush_thread = None
+        self.chunk_id = 0
+
+        # Deduplication state
+        self.last_line = ""
+        self.last_time = 0.0
+        self.last_progress_line = ""  # Track progress sequence key for deduplication
+        self.last_was_progress = False  # Track if last line was a progress bar
+
+    def start_capture(self):
+        """Start capturing console output and redirect stdout/stderr.
+
+        Notes:
+            In DDP training, only activates on rank 0/-1 to prevent duplicate logging.
+        """
+        if self.active or RANK not in {-1, 0}:
+            return
+
+        self.active = True
+        sys.stdout = self._ConsoleCapture(self.original_stdout, self._queue_log)
+        sys.stderr = self._ConsoleCapture(self.original_stderr, self._queue_log)
+
+        # Hook Ultralytics logger
+        try:
+            self._log_handler = self._LogHandler(self._queue_log)
+            logging.getLogger("ultralytics").addHandler(self._log_handler)
+        except Exception:
+            pass
+
+        # Start background flush thread for batched mode
+        if self.batch_size > 1:
+            self.flush_thread = threading.Thread(target=self._flush_worker, daemon=True)
+            self.flush_thread.start()
+
+    def stop_capture(self):
+        """Stop capturing console output and flush remaining buffer."""
+        if not self.active:
+            return
+
+        self.active = False
+        sys.stdout = self.original_stdout
+        sys.stderr = self.original_stderr
+
+        # Remove logging handler to prevent memory leak
+        if self._log_handler:
+            try:
+                logging.getLogger("ultralytics").removeHandler(self._log_handler)
+            except Exception:
+                pass
+            self._log_handler = None
+
+        # Final flush
+        self._flush_buffer()
+
+    def _queue_log(self, text):
+        """Queue console text with deduplication and timestamp processing."""
+        if not self.active:
+            return
+
+        current_time = time.time()
+
+        # Handle carriage returns and process lines
+        if "\r" in text:
+            text = text.split("\r")[-1]
+
+        lines = text.split("\n")
+        if lines and lines[-1] == "":
+            lines.pop()
+
+        for line in lines:
+            line = line.rstrip()
+
+            # Skip lines with only thin progress bars (partial progress)
+            if "─" in line:  # Has thin lines but no thick lines
+                continue
+
+            # Only show 100% completion lines for progress bars
+            if " ━━" in line:
+                is_complete = "100%" in line
+
+                # Skip ALL non-complete progress lines
+                if not is_complete:
+                    continue
+
+                # Extract sequence key to deduplicate multiple 100% lines for same sequence
+                parts = line.split()
+                seq_key = ""
+                if parts:
+                    # Check for epoch pattern (X/Y at start)
+                    if "/" in parts[0] and parts[0].replace("/", "").isdigit():
+                        seq_key = parts[0]  # e.g., "1/3"
+                    elif parts[0] == "Class" and len(parts) > 1:
+                        seq_key = f"{parts[0]}_{parts[1]}"  # e.g., "Class_train:" or "Class_val:"
+                    elif parts[0] in ("train:", "val:"):
+                        seq_key = parts[0]  # Phase identifier
+
+                # Skip if we already showed 100% for this sequence
+                if seq_key and self.last_progress_line == f"{seq_key}:done":
+                    continue
+
+                # Mark this sequence as done
+                if seq_key:
+                    self.last_progress_line = f"{seq_key}:done"
+
+                self.last_was_progress = True
+            else:
+                # Skip empty line after progress bar
+                if not line and self.last_was_progress:
+                    self.last_was_progress = False
+                    continue
+                self.last_was_progress = False
+
+            # General deduplication
+            if line == self.last_line and current_time - self.last_time < 0.1:
+                continue
+
+            self.last_line = line
+            self.last_time = current_time
+
+            # Add timestamp if needed
+            if not line.startswith("[20"):
+                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                line = f"[{timestamp}] {line}"
+
+            # Add to buffer and check if flush needed
+            should_flush = False
+            with self.buffer_lock:
+                self.buffer.append(line)
+                if len(self.buffer) >= self.batch_size:
+                    should_flush = True
+
+            # Flush outside lock to avoid deadlock
+            if should_flush:
+                self._flush_buffer()
+
+    def _flush_worker(self):
+        """Background worker that flushes buffer periodically."""
+        while self.active:
+            time.sleep(self.flush_interval)
+            if self.active:
+                self._flush_buffer()
+
+    def _flush_buffer(self):
+        """Flush buffered lines to destination and/or callback."""
+        with self.buffer_lock:
+            if not self.buffer:
+                return
+            lines = self.buffer.copy()
+            self.buffer.clear()
+            self.chunk_id += 1
+            chunk_id = self.chunk_id  # Capture under lock to avoid race
+
+        content = "\n".join(lines)
+        line_count = len(lines)
+
+        # Call custom callback if provided
+        if self.on_flush:
+            try:
+                self.on_flush(content, line_count, chunk_id)
+            except Exception:
+                pass  # Silently ignore callback errors to avoid flooding stderr
+
+        # Write to destination (file or API)
+        if self.destination is not None:
+            self._write_destination(content)
+
+    def _write_destination(self, content):
+        """Write content to file or API destination."""
+        try:
+            if self.is_api:
+                import requests
+
+                payload = {"timestamp": datetime.now().isoformat(), "message": content}
+                requests.post(str(self.destination), json=payload, timeout=5)
+            else:
+                self.destination.parent.mkdir(parents=True, exist_ok=True)
+                with self.destination.open("a", encoding="utf-8") as f:
+                    f.write(content + "\n")
+        except Exception as e:
+            print(f"Console logger write error: {e}", file=self.original_stderr)
+
+    class _ConsoleCapture:
+        """Lightweight stdout/stderr capture."""
+
+        __slots__ = ("callback", "original")
+
+        def __init__(self, original, callback):
+            """Initialize a stream wrapper that redirects writes to a callback while preserving the original."""
+            self.original = original
+            self.callback = callback
+
+        def write(self, text):
+            """Write text to the original stream and forward it to the capture callback."""
+            self.original.write(text)
+            self.callback(text)
+
+        def flush(self):
+            """Flush the wrapped stream to propagate buffered output promptly during console capture."""
+            self.original.flush()
+
+    class _LogHandler(logging.Handler):
+        """Lightweight logging handler."""
+
+        __slots__ = ("callback",)
+
+        def __init__(self, callback):
+            """Initialize a lightweight logging.Handler that forwards log records to the provided callback."""
+            super().__init__()
+            self.callback = callback
+
+        def emit(self, record):
+            """Format and forward LogRecord messages to the capture callback for unified log streaming."""
+            self.callback(self.format(record) + "\n")
+
+
+class SystemLogger:
+    """Log dynamic system metrics for training monitoring.
+
+    Captures real-time system metrics including CPU, RAM, disk I/O, network I/O, and NVIDIA GPU statistics for training
+    performance monitoring and analysis.
+
+    Attributes:
+        pynvml: NVIDIA pynvml module instance if successfully imported, None otherwise.
+        nvidia_initialized (bool): Whether NVIDIA GPU monitoring is available and initialized.
+        net_start: Initial network I/O counters for calculating cumulative usage.
+        disk_start: Initial disk I/O counters for calculating cumulative usage.
+
+    Examples:
+        Basic usage:
+        >>> logger = SystemLogger()
+        >>> metrics = logger.get_metrics()
+        >>> print(f"CPU: {metrics['cpu']}%, RAM: {metrics['ram']}%")
+        >>> if metrics["gpus"]:
+        ...     gpu0 = metrics["gpus"]["0"]
+        ...     print(f"GPU0: {gpu0['usage']}% usage, {gpu0['temp']}°C")
+
+        Training loop integration:
+        >>> system_logger = SystemLogger()
+        >>> for epoch in range(epochs):
+        ...     # Training code here
+        ...     metrics = system_logger.get_metrics()
+        ...     # Log to database/file
+    """
+
+    def __init__(self):
+        """Initialize the system logger."""
+        import psutil  # scoped as slow import
+
+        self.pynvml = None
+        self.nvidia_initialized = self._init_nvidia()
+        self.net_start = psutil.net_io_counters()
+        self.disk_start = psutil.disk_io_counters()
+
+        # For rate calculation
+        self._prev_net = self.net_start
+        self._prev_disk = self.disk_start
+        self._prev_time = time.time()
+
+    def _init_nvidia(self):
+        """Initialize NVIDIA GPU monitoring with pynvml."""
+        if MACOS:
+            return False
+
+        try:
+            check_requirements("nvidia-ml-py>=12.0.0")
+            self.pynvml = __import__("pynvml")
+            self.pynvml.nvmlInit()
+            return True
+        except Exception as e:
+            import torch
+
+            if torch.cuda.is_available():
+                LOGGER.warning(f"SystemLogger NVML init failed: {e}")
+            return False
+
+    def get_metrics(self, rates=False):
+        """Get current system metrics including CPU, RAM, disk, network, and GPU usage.
+
+        Collects comprehensive system metrics including CPU usage, RAM usage, disk I/O statistics, network I/O
+        statistics, and GPU metrics (if available).
+
+        Example output (rates=False, default):
+        ```python
+        {
+            "cpu": 45.2,
+            "ram": 78.9,
+            "disk": {"read_mb": 156.7, "write_mb": 89.3, "used_gb": 256.8},
+            "network": {"recv_mb": 157.2, "sent_mb": 89.1},
+            "gpus": {
+                "0": {"usage": 95.6, "memory": 85.4, "temp": 72, "power": 285},
+                "1": {"usage": 94.1, "memory": 82.7, "temp": 70, "power": 278},
+            },
+        }
+        ```
+
+        Example output (rates=True):
+        ```python
+        {
+            "cpu": 45.2,
+            "ram": 78.9,
+            "disk": {"read_mbs": 12.5, "write_mbs": 8.3, "used_gb": 256.8},
+            "network": {"recv_mbs": 5.2, "sent_mbs": 1.1},
+            "gpus": {
+                "0": {"usage": 95.6, "memory": 85.4, "temp": 72, "power": 285},
+            },
+        }
+        ```
+
+        Args:
+            rates (bool): If True, return disk/network as MB/s rates instead of cumulative MB.
+
+        Returns:
+            (dict): Metrics dictionary with cpu, ram, disk, network, and gpus keys.
+
+        Examples:
+            >>> logger = SystemLogger()
+            >>> logger.get_metrics()["cpu"]  # CPU percentage
+            >>> logger.get_metrics(rates=True)["network"]["recv_mbs"]  # MB/s download rate
+        """
+        import psutil  # scoped as slow import
+
+        net = psutil.net_io_counters()
+        disk = psutil.disk_io_counters()
+        memory = psutil.virtual_memory()
+        disk_usage = shutil.disk_usage("/")
+        now = time.time()
+
+        metrics = {
+            "cpu": round(psutil.cpu_percent(), 3),
+            "ram": round(memory.percent, 3),
+            "gpus": {},
+        }
+
+        # Calculate elapsed time since last call
+        elapsed = max(0.1, now - self._prev_time)  # Avoid division by zero
+
+        if rates:
+            # Calculate MB/s rates from delta since last call
+            metrics["disk"] = {
+                "read_mbs": round(max(0, (disk.read_bytes - self._prev_disk.read_bytes) / (1 << 20) / elapsed), 3),
+                "write_mbs": round(max(0, (disk.write_bytes - self._prev_disk.write_bytes) / (1 << 20) / elapsed), 3),
+                "used_gb": round(disk_usage.used / (1 << 30), 3),
+            }
+            metrics["network"] = {
+                "recv_mbs": round(max(0, (net.bytes_recv - self._prev_net.bytes_recv) / (1 << 20) / elapsed), 3),
+                "sent_mbs": round(max(0, (net.bytes_sent - self._prev_net.bytes_sent) / (1 << 20) / elapsed), 3),
+            }
+        else:
+            # Cumulative MB since initialization (original behavior)
+            metrics["disk"] = {
+                "read_mb": round((disk.read_bytes - self.disk_start.read_bytes) / (1 << 20), 3),
+                "write_mb": round((disk.write_bytes - self.disk_start.write_bytes) / (1 << 20), 3),
+                "used_gb": round(disk_usage.used / (1 << 30), 3),
+            }
+            metrics["network"] = {
+                "recv_mb": round((net.bytes_recv - self.net_start.bytes_recv) / (1 << 20), 3),
+                "sent_mb": round((net.bytes_sent - self.net_start.bytes_sent) / (1 << 20), 3),
+            }
+
+        # Always update previous values for accurate rate calculation on next call
+        self._prev_net = net
+        self._prev_disk = disk
+        self._prev_time = now
+
+        # Add GPU metrics (NVIDIA only)
+        if self.nvidia_initialized:
+            metrics["gpus"].update(self._get_nvidia_metrics())
+
+        return metrics
+
+    def _get_nvidia_metrics(self):
+        """Get NVIDIA GPU metrics including utilization, memory, temperature, and power."""
+        gpus = {}
+        if not self.nvidia_initialized or not self.pynvml:
+            return gpus
+        try:
+            device_count = self.pynvml.nvmlDeviceGetCount()
+            for i in range(device_count):
+                handle = self.pynvml.nvmlDeviceGetHandleByIndex(i)
+                util = self.pynvml.nvmlDeviceGetUtilizationRates(handle)
+                memory = self.pynvml.nvmlDeviceGetMemoryInfo(handle)
+                temp = self.pynvml.nvmlDeviceGetTemperature(handle, self.pynvml.NVML_TEMPERATURE_GPU)
+                power = self.pynvml.nvmlDeviceGetPowerUsage(handle) // 1000
+
+                gpus[str(i)] = {
+                    "usage": round(util.gpu, 3),
+                    "memory": round((memory.used / memory.total) * 100, 3),
+                    "temp": temp,
+                    "power": power,
+                }
+        except Exception:
+            pass
+        return gpus
+
+
+if __name__ == "__main__":
+    print("SystemLogger Real-time Metrics Monitor")
+    print("Press Ctrl+C to stop\n")
+
+    logger = SystemLogger()
+
+    try:
+        while True:
+            metrics = logger.get_metrics()
+
+            # Clear screen (works on most terminals)
+            print("\033[H\033[J", end="")
+
+            # Display system metrics
+            print(f"CPU: {metrics['cpu']:5.1f}%")
+            print(f"RAM: {metrics['ram']:5.1f}%")
+            print(f"Disk Read: {metrics['disk']['read_mb']:8.1f} MB")
+            print(f"Disk Write: {metrics['disk']['write_mb']:7.1f} MB")
+            print(f"Disk Used: {metrics['disk']['used_gb']:8.1f} GB")
+            print(f"Net Recv: {metrics['network']['recv_mb']:9.1f} MB")
+            print(f"Net Sent: {metrics['network']['sent_mb']:9.1f} MB")
+
+            # Display GPU metrics if available
+            if metrics["gpus"]:
+                print("\nGPU Metrics:")
+                for gpu_id, gpu_data in metrics["gpus"].items():
+                    print(
+                        f"  GPU {gpu_id}: {gpu_data['usage']:3}% | "
+                        f"Mem: {gpu_data['memory']:5.1f}% | "
+                        f"Temp: {gpu_data['temp']:2}°C | "
+                        f"Power: {gpu_data['power']:3}W"
+                    )
+            else:
+                print("\nGPU: No NVIDIA GPUs detected")
+
+            time.sleep(1)
+
+    except KeyboardInterrupt:
+        print("\n\nStopped monitoring.")

ultralytics/utils/metrics_3d.py

@@ -0,0 +1,267 @@
+# Ultralytics AGPL-3.0 License - https://ultralytics.com/license
+
+"""3D detection metrics for monocular 3D object detection.
+
+Ported from yolov5-3d/utils/metrics.py. Provides depth error, orientation error,
+3D center error, UV error, and grouped aggregation for matched prediction-GT pairs.
+"""
+
+import math
+
+import numpy as np
+
+
+def compute_depth_error(pred_depth, gt_depth, eps=1e-7):
+    """Compute depth error metrics between predicted and ground truth depth.
+
+    Args:
+        pred_depth: Predicted depth values (N,).
+        gt_depth: Ground truth depth values (N,).
+        eps: Small value to avoid division by zero.
+
+    Returns:
+        Dict with abs_error, rel_error, rmse.
+    """
+    pred_depth = np.asarray(pred_depth, dtype=np.float64)
+    gt_depth = np.asarray(gt_depth, dtype=np.float64)
+
+    valid = np.isfinite(pred_depth) & np.isfinite(gt_depth) & (gt_depth > 0)
+    if not np.any(valid):
+        return {"abs_error": 0.0, "rel_error": 0.0, "rmse": 0.0}
+
+    p, g = pred_depth[valid], gt_depth[valid]
+    abs_err = np.abs(p - g)
+    return {
+        "abs_error": float(np.mean(abs_err)),
+        "rel_error": float(np.mean(abs_err / (g + eps))),
+        "rmse": float(np.sqrt(np.mean((p - g) ** 2))),
+    }
+
+
+def compute_orientation_error(pred_yaw, gt_yaw):
+    """Compute mean absolute orientation error in degrees with wrap-around handling.
+
+    Args:
+        pred_yaw: Predicted yaw angles in radians (N,).
+        gt_yaw: Ground truth yaw angles in radians (N,).
+
+    Returns:
+        Mean absolute orientation error in degrees.
+    """
+    pred_yaw = np.asarray(pred_yaw, dtype=np.float64)
+    gt_yaw = np.asarray(gt_yaw, dtype=np.float64)
+
+    valid = np.isfinite(pred_yaw) & np.isfinite(gt_yaw)
+    if not np.any(valid):
+        return 0.0
+
+    diff = np.abs(pred_yaw[valid] - gt_yaw[valid])
+    diff = np.minimum(diff, 2 * math.pi - diff)
+    return float(np.mean(np.degrees(diff)))
+
+
+def count_valid_orientation_pairs(pred_yaw, gt_yaw):
+    """Count valid orientation pairs with finite prediction and ground-truth yaw."""
+    pred_yaw = np.asarray(pred_yaw, dtype=np.float64)
+    gt_yaw = np.asarray(gt_yaw, dtype=np.float64)
+    return int(np.sum(np.isfinite(pred_yaw) & np.isfinite(gt_yaw)))
+
+
+def compute_visible_orientation_metrics(pred_direct_yaw, pred_edge_yaw, gt_yaw):
+    """Compute visible-face direct and edge orientation errors against GT yaw."""
+    gt_yaw = np.asarray(gt_yaw, dtype=np.float64)
+    pred_direct_yaw = np.asarray(pred_direct_yaw, dtype=np.float64)
+    pred_edge_yaw = np.asarray(pred_edge_yaw, dtype=np.float64)
+
+    return {
+        "direct_orient_visible": compute_orientation_error(pred_direct_yaw, gt_yaw),
+        "edge_orient_visible": compute_orientation_error(pred_edge_yaw, gt_yaw),
+        "_direct_orient_visible_matched": count_valid_orientation_pairs(pred_direct_yaw, gt_yaw),
+        "_edge_orient_visible_matched": count_valid_orientation_pairs(pred_edge_yaw, gt_yaw),
+    }
+
+
+def compute_3d_center_error(pred_center, gt_center):
+    """Compute mean Euclidean distance between predicted and GT 3D centers.
+
+    Args:
+        pred_center: Predicted 3D centers (N, 3).
+        gt_center: Ground truth 3D centers (N, 3).
+
+    Returns:
+        Mean Euclidean distance in meters.
+    """
+    pred_center = np.asarray(pred_center, dtype=np.float64)
+    gt_center = np.asarray(gt_center, dtype=np.float64)
+
+    valid = np.all(np.isfinite(pred_center), axis=1) & np.all(np.isfinite(gt_center), axis=1)
+    if not np.any(valid):
+        return 0.0
+
+    return float(np.mean(np.linalg.norm(pred_center[valid] - gt_center[valid], axis=1)))
+
+
+def compute_size_error(pred_dims, gt_dims):
+    """Compute mean absolute size error for L, H, W dimensions.
+
+    Args:
+        pred_dims: Predicted dimensions (N, 3) - [l, h, w].
+        gt_dims: Ground truth dimensions (N, 3) - [l, h, w].
+
+    Returns:
+        Mean absolute size error in meters.
+    """
+    pred_dims = np.asarray(pred_dims, dtype=np.float64)
+    gt_dims = np.asarray(gt_dims, dtype=np.float64)
+
+    valid = np.all(np.isfinite(pred_dims), axis=1) & np.all(np.isfinite(gt_dims), axis=1)
+    if not np.any(valid):
+        return 0.0
+
+    return float(np.mean(np.abs(pred_dims[valid] - gt_dims[valid])))
+
+
+def compute_uv_error(pred_uv, gt_uv):
+    """Compute mean per-coordinate L1 pixel error between predicted and GT UV coordinates.
+
+    This matches the training-time UV logging, which averages absolute U and V errors
+    instead of using Euclidean point distance.
+    """
+    pred_uv = np.asarray(pred_uv, dtype=np.float64)
+    gt_uv = np.asarray(gt_uv, dtype=np.float64)
+
+    valid = np.all(np.isfinite(pred_uv), axis=1) & np.all(np.isfinite(gt_uv), axis=1)
+    if not np.any(valid):
+        return 0.0
+
+    diff = np.abs(pred_uv[valid] - gt_uv[valid])
+    return float(np.mean(diff))
+
+
+def empty_3d_metrics(include_orient=True, include_size=True, include_uv=True, include_visible_orient=False):
+    """Return default 3D metrics used for logging when no matches are available."""
+    metrics = {
+        "depth_abs": 0.0,
+        "depth_rel": 0.0,
+        "depth_rmse": 0.0,
+        "center": 0.0,
+        "matched": 0,
+    }
+    if include_uv:
+        metrics["uv"] = 0.0
+    if include_orient:
+        metrics["orient"] = 0.0
+    if include_size:
+        metrics["size"] = 0.0
+    if include_visible_orient:
+        metrics["direct_orient_visible"] = 0.0
+        metrics["edge_orient_visible"] = 0.0
+    return metrics
+
+
+def aggregate_3d_metric_groups(stats_by_group):
+    """Aggregate grouped 3D metrics with matched-count weighting."""
+    aggregated = {}
+    for group, entries in stats_by_group.items():
+        if not entries:
+            aggregated[group] = empty_3d_metrics(
+                include_orient=group == "whole",
+                include_size=group == "whole",
+                include_visible_orient=group == "face",
+            )
+            continue
+
+        template = {k: 0.0 for k in entries[0] if k != "matched" and not k.startswith("_")}
+        template["matched"] = 0
+        total_matched = sum(entry["matched"] for entry in entries)
+        total_pos_matched = sum(entry.get("_pos_matched", entry["matched"]) for entry in entries)
+        if total_matched <= 0 and total_pos_matched <= 0:
+            aggregated[group] = template
+            continue
+
+        for key in template:
+            if key == "matched":
+                continue
+            if key in {"depth_abs", "depth_rel", "depth_rmse", "center", "uv"}:
+                weight_key = "_pos_matched"
+            elif key == "direct_orient_visible":
+                weight_key = "_direct_orient_visible_matched"
+            elif key == "edge_orient_visible":
+                weight_key = "_edge_orient_visible_matched"
+            else:
+                weight_key = "matched"
+            total_weight = sum(entry.get(weight_key, entry.get("matched", 0)) for entry in entries)
+            if total_weight <= 0:
+                template[key] = float("nan") if key in {"direct_orient_visible", "edge_orient_visible"} else 0.0
+                continue
+            weighted = sum(
+                entry[key] * entry.get(weight_key, entry.get("matched", 0))
+                for entry in entries
+                if entry.get(weight_key, entry.get("matched", 0)) > 0
+            )
+            template[key] = round(weighted / total_weight, 5)
+        template["matched"] = total_matched
+        aggregated[group] = template
+    return aggregated
+
+
+def compute_3d_metrics_for_matched(
+    pred_3d_attrs,
+    gt_3d_attrs,
+    include_orient=True,
+    include_size=True,
+    include_uv=False,
+    include_visible_orient=False,
+):
+    """Compute 3D metrics for pre-matched prediction-GT pairs.
+
+    Args:
+        pred_3d_attrs: Dict with keys:
+            - center: (N, 3) predicted 3D centers [x, y, z]
+            - depth: (N,) predicted z3d
+            - yaw: (N,) predicted rotation_y in radians
+            - edge_yaw: (N,) predicted visible-face yaw in radians (optional)
+            - dims: (N, 3) predicted [l, h, w]
+            - uv: (N, 2) predicted [u, v] in pixels (optional)
+        gt_3d_attrs: Dict with same keys for ground truth.
+        include_orient: Whether to compute orientation error.
+        include_size: Whether to compute size error.
+        include_uv: Whether to compute UV pixel error.
+        include_visible_orient: Whether to compute visible-face direct and edge orientation errors.
+
+    Returns:
+        Dict with aggregated metrics and matched count.
+    """
+    n = len(pred_3d_attrs.get("depth", []))
+    if n == 0:
+        return empty_3d_metrics(
+            include_orient=include_orient,
+            include_size=include_size,
+            include_uv=include_uv,
+            include_visible_orient=include_visible_orient,
+        )
+
+    depth_m = compute_depth_error(pred_3d_attrs["depth"], gt_3d_attrs["depth"])
+    center_m = compute_3d_center_error(pred_3d_attrs["center"], gt_3d_attrs["center"])
+    metrics = {
+        "depth_abs": depth_m["abs_error"],
+        "depth_rel": depth_m["rel_error"],
+        "depth_rmse": depth_m["rmse"],
+        "center": center_m,
+        "matched": n,
+    }
+    if include_uv:
+        metrics["uv"] = compute_uv_error(pred_3d_attrs["uv"], gt_3d_attrs["uv"])
+    if include_orient:
+        metrics["orient"] = compute_orientation_error(pred_3d_attrs["yaw"], gt_3d_attrs["yaw"])
+    if include_size:
+        metrics["size"] = compute_size_error(pred_3d_attrs["dims"], gt_3d_attrs["dims"])
+    if include_visible_orient:
+        metrics.update(
+            compute_visible_orientation_metrics(
+                pred_3d_attrs["yaw"],
+                pred_3d_attrs.get("edge_yaw", np.full(n, np.nan, dtype=np.float64)),
+                gt_3d_attrs["yaw"],
+            )
+        )
+    return metrics

ultralytics/utils/nms.py

@@ -0,0 +1,338 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import sys
+import time
+
+import torch
+
+from ultralytics.utils import LOGGER
+from ultralytics.utils.metrics import batch_probiou, box_iou
+from ultralytics.utils.ops import xywh2xyxy
+
+
+def non_max_suppression(
+    prediction,
+    conf_thres: float = 0.25,
+    iou_thres: float = 0.45,
+    classes=None,
+    agnostic: bool = False,
+    multi_label: bool = False,
+    labels=(),
+    max_det: int = 300,
+    nc: int = 0,  # number of classes (optional)
+    max_time_img: float = 0.05,
+    max_nms: int = 30000,
+    max_wh: int = 7680,
+    rotated: bool = False,
+    end2end: bool = False,
+    return_idxs: bool = False,
+):
+    """Perform non-maximum suppression (NMS) on prediction results.
+
+    Applies NMS to filter overlapping bounding boxes based on confidence and IoU thresholds. Supports multiple detection
+    formats including standard boxes, rotated boxes, and masks.
+
+    Args:
+        prediction (torch.Tensor): Predictions with shape (batch_size, num_classes + 4 + num_masks, num_boxes)
+            containing boxes, classes, and optional masks.
+        conf_thres (float): Confidence threshold for filtering detections. Valid values are between 0.0 and 1.0.
+        iou_thres (float): IoU threshold for NMS filtering. Valid values are between 0.0 and 1.0.
+        classes (list[int], optional): List of class indices to consider. If None, all classes are considered.
+        agnostic (bool): Whether to perform class-agnostic NMS.
+        multi_label (bool): Whether each box can have multiple labels.
+        labels (list[torch.Tensor]): A priori labels for each image.
+        max_det (int): Maximum number of detections to keep per image.
+        nc (int): Number of classes. Indices after this are considered masks.
+        max_time_img (float): Maximum time in seconds for processing one image.
+        max_nms (int): Maximum number of boxes for NMS.
+        max_wh (int): Maximum box width and height in pixels.
+        rotated (bool): Whether to handle Oriented Bounding Boxes (OBB).
+        end2end (bool): Whether the model is end-to-end and doesn't require NMS.
+        return_idxs (bool): Whether to return the indices of kept detections.
+
+    Returns:
+        (list[torch.Tensor] | tuple[list[torch.Tensor], list[torch.Tensor]]): List of detections per image with shape
+            (num_boxes, 6 + num_masks) containing (x1, y1, x2, y2, confidence, class, mask1, mask2, ...). If
+            return_idxs=True, returns a tuple of (output, keepi) where keepi contains indices of kept detections.
+    """
+    # Checks
+    assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
+    assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"
+    if isinstance(prediction, (list, tuple)):  # YOLOv8 model in validation model, output = (inference_out, loss_out)
+        prediction = prediction[0]  # select only inference output
+    if classes is not None:
+        classes = torch.tensor(classes, device=prediction.device)
+
+    if prediction.shape[-1] == 6 or end2end:  # end-to-end model (BNC, i.e. 1,300,6)
+        output = [pred[pred[:, 4] > conf_thres][:max_det] for pred in prediction]
+        if classes is not None:
+            output = [pred[(pred[:, 5:6] == classes).any(1)] for pred in output]
+        return output
+
+    bs = prediction.shape[0]  # batch size (BCN, i.e. 1,84,6300)
+    nc = nc or (prediction.shape[1] - 4)  # number of classes
+    extra = prediction.shape[1] - nc - 4  # number of extra info
+    mi = 4 + nc  # mask start index
+    xc = prediction[:, 4:mi].amax(1) > conf_thres  # candidates
+    xinds = torch.arange(prediction.shape[-1], device=prediction.device).expand(bs, -1)[..., None]  # to track idxs
+
+    # Settings
+    # min_wh = 2  # (pixels) minimum box width and height
+    time_limit = 2.0 + max_time_img * bs  # seconds to quit after
+    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
+
+    prediction = prediction.transpose(-1, -2)  # shape(1,84,6300) to shape(1,6300,84)
+    if not rotated:
+        prediction[..., :4] = xywh2xyxy(prediction[..., :4])  # xywh to xyxy
+
+    t = time.time()
+    output = [torch.zeros((0, 6 + extra), device=prediction.device)] * bs
+    keepi = [torch.zeros((0, 1), device=prediction.device)] * bs  # to store the kept idxs
+    for xi, (x, xk) in enumerate(zip(prediction, xinds)):  # image index, (preds, preds indices)
+        # Apply constraints
+        # x[((x[:, 2:4] < min_wh) | (x[:, 2:4] > max_wh)).any(1), 4] = 0  # width-height
+        filt = xc[xi]  # confidence
+        x = x[filt]
+        if return_idxs:
+            xk = xk[filt]
+
+        # Cat apriori labels if autolabelling
+        if labels and len(labels[xi]) and not rotated:
+            lb = labels[xi]
+            v = torch.zeros((len(lb), nc + extra + 4), device=x.device)
+            v[:, :4] = xywh2xyxy(lb[:, 1:5])  # box
+            v[range(len(lb)), lb[:, 0].long() + 4] = 1.0  # cls
+            x = torch.cat((x, v), 0)
+
+        # If none remain process next image
+        if not x.shape[0]:
+            continue
+
+        # Detections matrix nx6 (xyxy, conf, cls)
+        box, cls, mask = x.split((4, nc, extra), 1)
+
+        if multi_label:
+            i, j = torch.where(cls > conf_thres)
+            x = torch.cat((box[i], x[i, 4 + j, None], j[:, None].float(), mask[i]), 1)
+            if return_idxs:
+                xk = xk[i]
+        else:  # best class only
+            conf, j = cls.max(1, keepdim=True)
+            filt = conf.view(-1) > conf_thres
+            x = torch.cat((box, conf, j.float(), mask), 1)[filt]
+            if return_idxs:
+                xk = xk[filt]
+
+        # Filter by class
+        if classes is not None:
+            filt = (x[:, 5:6] == classes).any(1)
+            x = x[filt]
+            if return_idxs:
+                xk = xk[filt]
+
+        # Check shape
+        n = x.shape[0]  # number of boxes
+        if not n:  # no boxes
+            continue
+        if n > max_nms:  # excess boxes
+            filt = x[:, 4].argsort(descending=True)[:max_nms]  # sort by confidence and remove excess boxes
+            x = x[filt]
+            if return_idxs:
+                xk = xk[filt]
+
+        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
+        scores = x[:, 4]  # scores
+        if rotated:
+            boxes = torch.cat((x[:, :2] + c, x[:, 2:4], x[:, -1:]), dim=-1)  # xywhr
+            i = TorchNMS.fast_nms(boxes, scores, iou_thres, iou_func=batch_probiou)
+        else:
+            boxes = x[:, :4] + c  # boxes (offset by class)
+            # Speed strategy: torchvision for val or already loaded (faster), TorchNMS for predict (lower latency)
+            if "torchvision" in sys.modules:
+                import torchvision  # scope as slow import
+
+                i = torchvision.ops.nms(boxes, scores, iou_thres)
+            else:
+                i = TorchNMS.nms(boxes, scores, iou_thres)
+        i = i[:max_det]  # limit detections
+
+        output[xi] = x[i]
+        if return_idxs:
+            keepi[xi] = xk[i].view(-1)
+        if (time.time() - t) > time_limit:
+            LOGGER.warning(f"NMS time limit {time_limit:.3f}s exceeded")
+            break  # time limit exceeded
+
+    return (output, keepi) if return_idxs else output
+
+
+class TorchNMS:
+    """Ultralytics custom NMS implementation optimized for YOLO.
+
+    This class provides static methods for performing non-maximum suppression (NMS) operations on bounding boxes,
+    including standard NMS, fast NMS, and batched NMS for multi-class scenarios.
+
+    Methods:
+        fast_nms: Fast-NMS using upper triangular matrix operations.
+        nms: Optimized NMS with early termination that matches torchvision behavior exactly.
+        batched_nms: Batched NMS for class-aware suppression.
+
+    Examples:
+        Perform standard NMS on boxes and scores
+        >>> boxes = torch.tensor([[0, 0, 10, 10], [5, 5, 15, 15]])
+        >>> scores = torch.tensor([0.9, 0.8])
+        >>> keep = TorchNMS.nms(boxes, scores, 0.5)
+    """
+
+    @staticmethod
+    def fast_nms(
+        boxes: torch.Tensor,
+        scores: torch.Tensor,
+        iou_threshold: float,
+        use_triu: bool = True,
+        iou_func=box_iou,
+        exit_early: bool = True,
+    ) -> torch.Tensor:
+        """Fast-NMS implementation from https://arxiv.org/pdf/1904.02689 using upper triangular matrix operations.
+
+        Args:
+            boxes (torch.Tensor): Bounding boxes with shape (N, 4) in xyxy format.
+            scores (torch.Tensor): Confidence scores with shape (N,).
+            iou_threshold (float): IoU threshold for suppression.
+            use_triu (bool): Whether to use torch.triu operator for upper triangular matrix operations.
+            iou_func (callable): Function to compute IoU between boxes.
+            exit_early (bool): Whether to exit early if there are no boxes.
+
+        Returns:
+            (torch.Tensor): Indices of boxes to keep after NMS.
+
+        Examples:
+            Apply NMS to a set of boxes
+            >>> boxes = torch.tensor([[0, 0, 10, 10], [5, 5, 15, 15]])
+            >>> scores = torch.tensor([0.9, 0.8])
+            >>> keep = TorchNMS.fast_nms(boxes, scores, 0.5)
+        """
+        if boxes.numel() == 0 and exit_early:
+            return torch.empty((0,), dtype=torch.int64, device=boxes.device)
+
+        sorted_idx = torch.argsort(scores, descending=True)
+        boxes = boxes[sorted_idx]
+        ious = iou_func(boxes, boxes)
+        if use_triu:
+            ious = ious.triu_(diagonal=1)
+            # NOTE: handle the case when len(boxes) hence exportable by eliminating if-else condition
+            pick = torch.nonzero((ious >= iou_threshold).sum(0) <= 0).squeeze_(-1)
+        else:
+            n = boxes.shape[0]
+            row_idx = torch.arange(n, device=boxes.device).view(-1, 1).expand(-1, n)
+            col_idx = torch.arange(n, device=boxes.device).view(1, -1).expand(n, -1)
+            upper_mask = row_idx < col_idx
+            ious = ious * upper_mask
+            # Zeroing these scores ensures the additional indices would not affect the final results
+            scores_ = scores[sorted_idx]
+            scores_[~((ious >= iou_threshold).sum(0) <= 0)] = 0
+            scores[sorted_idx] = scores_  # update original tensor for NMSModel
+            # NOTE: return indices with fixed length to avoid TFLite reshape error
+            pick = torch.topk(scores_, scores_.shape[0]).indices
+        return sorted_idx[pick]
+
+    @staticmethod
+    def nms(boxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float) -> torch.Tensor:
+        """Optimized NMS with early termination that matches torchvision behavior exactly.
+
+        Args:
+            boxes (torch.Tensor): Bounding boxes with shape (N, 4) in xyxy format.
+            scores (torch.Tensor): Confidence scores with shape (N,).
+            iou_threshold (float): IoU threshold for suppression.
+
+        Returns:
+            (torch.Tensor): Indices of boxes to keep after NMS.
+
+        Examples:
+            Apply NMS to a set of boxes
+            >>> boxes = torch.tensor([[0, 0, 10, 10], [5, 5, 15, 15]])
+            >>> scores = torch.tensor([0.9, 0.8])
+            >>> keep = TorchNMS.nms(boxes, scores, 0.5)
+        """
+        if boxes.numel() == 0:
+            return torch.empty((0,), dtype=torch.int64, device=boxes.device)
+
+        # Pre-allocate and extract coordinates once
+        x1, y1, x2, y2 = boxes.unbind(1)
+        areas = (x2 - x1) * (y2 - y1)
+
+        # Sort by scores descending
+        order = scores.argsort(0, descending=True)
+
+        # Pre-allocate keep list with maximum possible size
+        keep = torch.zeros(order.numel(), dtype=torch.int64, device=boxes.device)
+        keep_idx = 0
+        while order.numel() > 0:
+            i = order[0]
+            keep[keep_idx] = i
+            keep_idx += 1
+
+            if order.numel() == 1:
+                break
+            # Vectorized IoU calculation for remaining boxes
+            rest = order[1:]
+            xx1 = torch.maximum(x1[i], x1[rest])
+            yy1 = torch.maximum(y1[i], y1[rest])
+            xx2 = torch.minimum(x2[i], x2[rest])
+            yy2 = torch.minimum(y2[i], y2[rest])
+
+            # Fast intersection and IoU
+            w = (xx2 - xx1).clamp_(min=0)
+            h = (yy2 - yy1).clamp_(min=0)
+            inter = w * h
+            # Early exit: skip IoU calculation if no intersection
+            if inter.sum() == 0:
+                # No overlaps with current box, keep all remaining boxes
+                order = rest
+                continue
+            iou = inter / (areas[i] + areas[rest] - inter)
+            # Keep boxes with IoU <= threshold
+            order = rest[iou <= iou_threshold]
+
+        return keep[:keep_idx]
+
+    @staticmethod
+    def batched_nms(
+        boxes: torch.Tensor,
+        scores: torch.Tensor,
+        idxs: torch.Tensor,
+        iou_threshold: float,
+        use_fast_nms: bool = False,
+    ) -> torch.Tensor:
+        """Batched NMS for class-aware suppression.
+
+        Args:
+            boxes (torch.Tensor): Bounding boxes with shape (N, 4) in xyxy format.
+            scores (torch.Tensor): Confidence scores with shape (N,).
+            idxs (torch.Tensor): Class indices with shape (N,).
+            iou_threshold (float): IoU threshold for suppression.
+            use_fast_nms (bool): Whether to use the Fast-NMS implementation.
+
+        Returns:
+            (torch.Tensor): Indices of boxes to keep after NMS.
+
+        Examples:
+            Apply batched NMS across multiple classes
+            >>> boxes = torch.tensor([[0, 0, 10, 10], [5, 5, 15, 15]])
+            >>> scores = torch.tensor([0.9, 0.8])
+            >>> idxs = torch.tensor([0, 1])
+            >>> keep = TorchNMS.batched_nms(boxes, scores, idxs, 0.5)
+        """
+        if boxes.numel() == 0:
+            return torch.empty((0,), dtype=torch.int64, device=boxes.device)
+
+        # Strategy: offset boxes by class index to prevent cross-class suppression
+        max_coordinate = boxes.max()
+        offsets = idxs.to(boxes) * (max_coordinate + 1)
+        boxes_for_nms = boxes + offsets[:, None]
+
+        return (
+            TorchNMS.fast_nms(boxes_for_nms, scores, iou_threshold)
+            if use_fast_nms
+            else TorchNMS.nms(boxes_for_nms, scores, iou_threshold)
+        )

ultralytics/utils/ops.py

@@ -0,0 +1,672 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import contextlib
+import math
+import re
+import time
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from ultralytics.utils import NOT_MACOS14
+
+
+class Profile(contextlib.ContextDecorator):
+    """Ultralytics Profile class for timing code execution.
+
+    Use as a decorator with @Profile() or as a context manager with 'with Profile():'. Provides accurate timing
+    measurements with CUDA synchronization support for GPU operations.
+
+    Attributes:
+        t (float): Accumulated time in seconds.
+        device (torch.device): Device used for model inference.
+        cuda (bool): Whether CUDA is being used for timing synchronization.
+
+    Examples:
+        Use as a context manager to time code execution
+        >>> with Profile(device=device) as dt:
+        ...     pass  # slow operation here
+        >>> print(dt)  # prints "Elapsed time is 9.5367431640625e-07 s"
+
+        Use as a decorator to time function execution
+        >>> @Profile()
+        ... def slow_function():
+        ...     time.sleep(0.1)
+    """
+
+    def __init__(self, t: float = 0.0, device: torch.device | None = None):
+        """Initialize the Profile class.
+
+        Args:
+            t (float): Initial accumulated time in seconds.
+            device (torch.device, optional): Device used for model inference to enable CUDA synchronization.
+        """
+        self.t = t
+        self.device = device
+        self.cuda = bool(device and str(device).startswith("cuda"))
+
+    def __enter__(self):
+        """Start timing."""
+        self.start = self.time()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        """Stop timing."""
+        self.dt = self.time() - self.start  # delta-time
+        self.t += self.dt  # accumulate dt
+
+    def __str__(self):
+        """Return a human-readable string representing the accumulated elapsed time."""
+        return f"Elapsed time is {self.t} s"
+
+    def time(self):
+        """Get current time with CUDA synchronization if applicable."""
+        if self.cuda:
+            torch.cuda.synchronize(self.device)
+        return time.perf_counter()
+
+
+def segment2box(segment, width: int = 640, height: int = 640):
+    """Convert segment coordinates to bounding box coordinates.
+
+    Converts a single segment label to a box label by finding the minimum and maximum x and y coordinates. Applies
+    inside-image constraint and clips coordinates when necessary.
+
+    Args:
+        segment (np.ndarray): Segment coordinates in format (N, 2) where N is number of points.
+        width (int): Width of the image in pixels.
+        height (int): Height of the image in pixels.
+
+    Returns:
+        (np.ndarray): Bounding box coordinates in xyxy format [x1, y1, x2, y2].
+    """
+    x, y = segment.T  # segment xy
+    # Clip coordinates if 3 out of 4 sides are outside the image
+    if np.array([x.min() < 0, y.min() < 0, x.max() > width, y.max() > height]).sum() >= 3:
+        x = x.clip(0, width)
+        y = y.clip(0, height)
+    inside = (x > 0) & (y > 0) & (x < width) & (y < height)
+    x = x[inside]
+    y = y[inside]
+    return (
+        np.array([x.min(), y.min(), x.max(), y.max()], dtype=segment.dtype)
+        if any(x)
+        else np.zeros(4, dtype=segment.dtype)
+    )  # xyxy
+
+
+def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None, padding: bool = True, xywh: bool = False):
+    """Rescale bounding boxes from one image shape to another.
+
+    Rescales bounding boxes from img1_shape to img0_shape, accounting for padding and aspect ratio changes. Supports
+    both xyxy and xywh box formats.
+
+    Args:
+        img1_shape (tuple): Shape of the source image (height, width).
+        boxes (torch.Tensor): Bounding boxes to rescale in format (N, 4).
+        img0_shape (tuple): Shape of the target image (height, width).
+        ratio_pad (tuple, optional): Tuple of (ratio, pad) for scaling. If None, calculated from image shapes.
+        padding (bool): Whether boxes are based on YOLO-style augmented images with padding.
+        xywh (bool): Whether box format is xywh (True) or xyxy (False).
+
+    Returns:
+        (torch.Tensor): Rescaled bounding boxes in the same format as input.
+    """
+    if ratio_pad is None:  # calculate from img0_shape
+        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+        pad_x = round((img1_shape[1] - img0_shape[1] * gain) / 2 - 0.1)
+        pad_y = round((img1_shape[0] - img0_shape[0] * gain) / 2 - 0.1)
+    else:
+        gain = ratio_pad[0][0]
+        pad_x, pad_y = ratio_pad[1]
+
+    if padding:
+        boxes[..., 0] -= pad_x  # x padding
+        boxes[..., 1] -= pad_y  # y padding
+        if not xywh:
+            boxes[..., 2] -= pad_x  # x padding
+            boxes[..., 3] -= pad_y  # y padding
+    boxes[..., :4] /= gain
+    return boxes if xywh else clip_boxes(boxes, img0_shape)
+
+
+def make_divisible(x: int, divisor):
+    """Return the nearest number that is divisible by the given divisor.
+
+    Args:
+        x (int): The number to make divisible.
+        divisor (int | torch.Tensor): The divisor.
+
+    Returns:
+        (int): The nearest number divisible by the divisor.
+    """
+    if isinstance(divisor, torch.Tensor):
+        divisor = int(divisor.max())  # to int
+    return math.ceil(x / divisor) * divisor
+
+
+def clip_boxes(boxes, shape):
+    """Clip bounding boxes to image boundaries.
+
+    Args:
+        boxes (torch.Tensor | np.ndarray): Bounding boxes to clip.
+        shape (tuple): Image shape as HWC or HW (supports both).
+
+    Returns:
+        (torch.Tensor | np.ndarray): Clipped bounding boxes.
+    """
+    h, w = shape[:2]  # supports both HWC or HW shapes
+    if isinstance(boxes, torch.Tensor):  # faster individually
+        if NOT_MACOS14:
+            boxes[..., 0].clamp_(0, w)  # x1
+            boxes[..., 1].clamp_(0, h)  # y1
+            boxes[..., 2].clamp_(0, w)  # x2
+            boxes[..., 3].clamp_(0, h)  # y2
+        else:  # Apple macOS14 MPS bug https://github.com/ultralytics/ultralytics/pull/21878
+            boxes[..., 0] = boxes[..., 0].clamp(0, w)
+            boxes[..., 1] = boxes[..., 1].clamp(0, h)
+            boxes[..., 2] = boxes[..., 2].clamp(0, w)
+            boxes[..., 3] = boxes[..., 3].clamp(0, h)
+    else:  # np.array (faster grouped)
+        boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(0, w)  # x1, x2
+        boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(0, h)  # y1, y2
+    return boxes
+
+
+def clip_coords(coords, shape):
+    """Clip line coordinates to image boundaries.
+
+    Args:
+        coords (torch.Tensor | np.ndarray): Line coordinates to clip.
+        shape (tuple): Image shape as HWC or HW (supports both).
+
+    Returns:
+        (torch.Tensor | np.ndarray): Clipped coordinates.
+    """
+    h, w = shape[:2]  # supports both HWC or HW shapes
+    if isinstance(coords, torch.Tensor):
+        if NOT_MACOS14:
+            coords[..., 0].clamp_(0, w)  # x
+            coords[..., 1].clamp_(0, h)  # y
+        else:  # Apple macOS14 MPS bug https://github.com/ultralytics/ultralytics/pull/21878
+            coords[..., 0] = coords[..., 0].clamp(0, w)
+            coords[..., 1] = coords[..., 1].clamp(0, h)
+    else:  # np.array
+        coords[..., 0] = coords[..., 0].clip(0, w)  # x
+        coords[..., 1] = coords[..., 1].clip(0, h)  # y
+    return coords
+
+
+def xyxy2xywh(x):
+    """Convert bounding box coordinates from (x1, y1, x2, y2) format to (x, y, width, height) format where (x1, y1) is
+    the top-left corner and (x2, y2) is the bottom-right corner.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Input bounding box coordinates in (x1, y1, x2, y2) format.
+
+    Returns:
+        (np.ndarray | torch.Tensor): Bounding box coordinates in (x, y, width, height) format.
+    """
+    assert x.shape[-1] == 4, f"input shape last dimension expected 4 but input shape is {x.shape}"
+    y = empty_like(x)  # faster than clone/copy
+    x1, y1, x2, y2 = x[..., 0], x[..., 1], x[..., 2], x[..., 3]
+    y[..., 0] = (x1 + x2) / 2  # x center
+    y[..., 1] = (y1 + y2) / 2  # y center
+    y[..., 2] = x2 - x1  # width
+    y[..., 3] = y2 - y1  # height
+    return y
+
+
+def xywh2xyxy(x):
+    """Convert bounding box coordinates from (x, y, width, height) format to (x1, y1, x2, y2) format where (x1, y1) is
+    the top-left corner and (x2, y2) is the bottom-right corner. Note: ops per 2 channels faster than per channel.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Input bounding box coordinates in (x, y, width, height) format.
+
+    Returns:
+        (np.ndarray | torch.Tensor): Bounding box coordinates in (x1, y1, x2, y2) format.
+    """
+    assert x.shape[-1] == 4, f"input shape last dimension expected 4 but input shape is {x.shape}"
+    y = empty_like(x)  # faster than clone/copy
+    xy = x[..., :2]  # centers
+    wh = x[..., 2:] / 2  # half width-height
+    y[..., :2] = xy - wh  # top left xy
+    y[..., 2:] = xy + wh  # bottom right xy
+    return y
+
+
+def xywhn2xyxy(x, w: int = 640, h: int = 640, padw: int = 0, padh: int = 0):
+    """Convert normalized bounding box coordinates to pixel coordinates.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Normalized bounding box coordinates in (x, y, w, h) format.
+        w (int): Image width in pixels.
+        h (int): Image height in pixels.
+        padw (int): Padding width in pixels.
+        padh (int): Padding height in pixels.
+
+    Returns:
+        (np.ndarray | torch.Tensor): Bounding box coordinates in (x1, y1, x2, y2) format.
+    """
+    assert x.shape[-1] == 4, f"input shape last dimension expected 4 but input shape is {x.shape}"
+    y = empty_like(x)  # faster than clone/copy
+    xc, yc, xw, xh = x[..., 0], x[..., 1], x[..., 2], x[..., 3]
+    half_w, half_h = xw / 2, xh / 2
+    y[..., 0] = w * (xc - half_w) + padw  # top left x
+    y[..., 1] = h * (yc - half_h) + padh  # top left y
+    y[..., 2] = w * (xc + half_w) + padw  # bottom right x
+    y[..., 3] = h * (yc + half_h) + padh  # bottom right y
+    return y
+
+
+def xyxy2xywhn(x, w: int = 640, h: int = 640, clip: bool = False, eps: float = 0.0):
+    """Convert bounding box coordinates from (x1, y1, x2, y2) format to normalized (x, y, width, height) format. x, y,
+    width and height are normalized to image dimensions.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Input bounding box coordinates in (x1, y1, x2, y2) format.
+        w (int): Image width in pixels.
+        h (int): Image height in pixels.
+        clip (bool): Whether to clip boxes to image boundaries.
+        eps (float): Minimum value for box width and height.
+
+    Returns:
+        (np.ndarray | torch.Tensor): Normalized bounding box coordinates in (x, y, width, height) format.
+    """
+    if clip:
+        x = clip_boxes(x, (h - eps, w - eps))
+    assert x.shape[-1] == 4, f"input shape last dimension expected 4 but input shape is {x.shape}"
+    y = empty_like(x)  # faster than clone/copy
+    x1, y1, x2, y2 = x[..., 0], x[..., 1], x[..., 2], x[..., 3]
+    y[..., 0] = ((x1 + x2) / 2) / w  # x center
+    y[..., 1] = ((y1 + y2) / 2) / h  # y center
+    y[..., 2] = (x2 - x1) / w  # width
+    y[..., 3] = (y2 - y1) / h  # height
+    return y
+
+
+def xywh2ltwh(x):
+    """Convert bounding box format from [x, y, w, h] to [x1, y1, w, h] where x1, y1 are top-left coordinates.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Input bounding box coordinates in xywh format.
+
+    Returns:
+        (np.ndarray | torch.Tensor): Bounding box coordinates in ltwh format.
+    """
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 0] = x[..., 0] - x[..., 2] / 2  # top left x
+    y[..., 1] = x[..., 1] - x[..., 3] / 2  # top left y
+    return y
+
+
+def xyxy2ltwh(x):
+    """Convert bounding boxes from [x1, y1, x2, y2] to [x1, y1, w, h] format.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Input bounding box coordinates in xyxy format.
+
+    Returns:
+        (np.ndarray | torch.Tensor): Bounding box coordinates in ltwh format.
+    """
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 2] = x[..., 2] - x[..., 0]  # width
+    y[..., 3] = x[..., 3] - x[..., 1]  # height
+    return y
+
+
+def ltwh2xywh(x):
+    """Convert bounding boxes from [x1, y1, w, h] to [x, y, w, h] where xy1=top-left, xy=center.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Input bounding box coordinates.
+
+    Returns:
+        (np.ndarray | torch.Tensor): Bounding box coordinates in xywh format.
+    """
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 0] = x[..., 0] + x[..., 2] / 2  # center x
+    y[..., 1] = x[..., 1] + x[..., 3] / 2  # center y
+    return y
+
+
+def xyxyxyxy2xywhr(x):
+    """Convert batched Oriented Bounding Boxes (OBB) from [xy1, xy2, xy3, xy4] to [xywh, rotation] format.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Input box corners with shape (N, 8) in [xy1, xy2, xy3, xy4] format.
+
+    Returns:
+        (np.ndarray | torch.Tensor): Converted data in [cx, cy, w, h, rotation] format with shape (N, 5). Rotation
+            values are in radians from [-pi/4, 3pi/4).
+    """
+    is_torch = isinstance(x, torch.Tensor)
+    points = x.cpu().numpy() if is_torch else x
+    points = points.reshape(len(x), -1, 2)
+    rboxes = []
+    for pts in points:
+        # NOTE: Use cv2.minAreaRect to get accurate xywhr,
+        # especially some objects are cut off by augmentations in dataloader.
+        (cx, cy), (w, h), angle = cv2.minAreaRect(pts)
+        # convert angle to radian and normalize to [-pi/4, 3pi/4)
+        theta = angle / 180 * np.pi
+        if w < h:
+            w, h = h, w
+            theta += np.pi / 2
+        while theta >= 3 * np.pi / 4:
+            theta -= np.pi
+        while theta < -np.pi / 4:
+            theta += np.pi
+        rboxes.append([cx, cy, w, h, theta])
+    return torch.tensor(rboxes, device=x.device, dtype=x.dtype) if is_torch else np.asarray(rboxes)
+
+
+def xywhr2xyxyxyxy(x):
+    """Convert batched Oriented Bounding Boxes (OBB) from [xywh, rotation] to [xy1, xy2, xy3, xy4] format.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Boxes in [cx, cy, w, h, rotation] format with shape (N, 5) or (B, N, 5). Rotation
+            values should be in radians from [-pi/4, 3pi/4).
+
+    Returns:
+        (np.ndarray | torch.Tensor): Converted corner points with shape (N, 4, 2) or (B, N, 4, 2).
+    """
+    cos, sin, cat, stack = (
+        (torch.cos, torch.sin, torch.cat, torch.stack)
+        if isinstance(x, torch.Tensor)
+        else (np.cos, np.sin, np.concatenate, np.stack)
+    )
+
+    ctr = x[..., :2]
+    w, h, angle = (x[..., i : i + 1] for i in range(2, 5))
+    cos_value, sin_value = cos(angle), sin(angle)
+    vec1 = [w / 2 * cos_value, w / 2 * sin_value]
+    vec2 = [-h / 2 * sin_value, h / 2 * cos_value]
+    vec1 = cat(vec1, -1)
+    vec2 = cat(vec2, -1)
+    pt1 = ctr + vec1 + vec2
+    pt2 = ctr + vec1 - vec2
+    pt3 = ctr - vec1 - vec2
+    pt4 = ctr - vec1 + vec2
+    return stack([pt1, pt2, pt3, pt4], -2)
+
+
+def ltwh2xyxy(x):
+    """Convert bounding box from [x1, y1, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right.
+
+    Args:
+        x (np.ndarray | torch.Tensor): Input bounding box coordinates.
+
+    Returns:
+        (np.ndarray | torch.Tensor): Bounding box coordinates in xyxy format.
+    """
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[..., 2] = x[..., 2] + x[..., 0]  # x2
+    y[..., 3] = x[..., 3] + x[..., 1]  # y2
+    return y
+
+
+def segments2boxes(segments):
+    """Convert segment coordinates to bounding box labels in xywh format.
+
+    Args:
+        segments (list): List of segments where each segment is a list of points, each point is [x, y] coordinates.
+
+    Returns:
+        (np.ndarray): Bounding box coordinates in xywh format.
+    """
+    boxes = []
+    for s in segments:
+        x, y = s.T  # segment xy
+        boxes.append([x.min(), y.min(), x.max(), y.max()])  # cls, xyxy
+    return xyxy2xywh(np.array(boxes))  # cls, xywh
+
+
+def resample_segments(segments, n: int = 1000):
+    """Resample segments to n points each using linear interpolation.
+
+    Args:
+        segments (list): List of (N, 2) arrays where N is the number of points in each segment.
+        n (int): Number of points to resample each segment to.
+
+    Returns:
+        (list): Resampled segments with n points each.
+    """
+    for i, s in enumerate(segments):
+        if len(s) == n:
+            continue
+        s = np.concatenate((s, s[0:1, :]), axis=0)
+        x = np.linspace(0, len(s) - 1, n - len(s) if len(s) < n else n)
+        xp = np.arange(len(s))
+        x = np.insert(x, np.searchsorted(x, xp), xp) if len(s) < n else x
+        segments[i] = (
+            np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)], dtype=np.float32).reshape(2, -1).T
+        )  # segment xy
+    return segments
+
+
+def crop_mask(masks: torch.Tensor, boxes: torch.Tensor) -> torch.Tensor:
+    """Crop masks to bounding box regions.
+
+    Args:
+        masks (torch.Tensor): Masks with shape (N, H, W).
+        boxes (torch.Tensor): Bounding box coordinates with shape (N, 4) in xyxy pixel format.
+
+    Returns:
+        (torch.Tensor): Cropped masks.
+    """
+    if boxes.device != masks.device:
+        boxes = boxes.to(masks.device)
+    n, h, w = masks.shape
+    if n < 50 and not masks.is_cuda:  # faster for fewer masks (predict)
+        for i, (x1, y1, x2, y2) in enumerate(boxes.round().int()):
+            masks[i, :y1] = 0
+            masks[i, y2:] = 0
+            masks[i, :, :x1] = 0
+            masks[i, :, x2:] = 0
+        return masks
+    else:  # faster for more masks (val)
+        x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1)  # x1 shape(n,1,1)
+        r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :]  # rows shape(1,1,w)
+        c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None]  # cols shape(1,h,1)
+        return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
+
+
+def process_mask(protos, masks_in, bboxes, shape, upsample: bool = False):
+    """Apply masks to bounding boxes using mask head output.
+
+    Args:
+        protos (torch.Tensor): Mask prototypes with shape (mask_dim, mask_h, mask_w).
+        masks_in (torch.Tensor): Mask coefficients with shape (N, mask_dim) where N is number of masks after NMS.
+        bboxes (torch.Tensor): Bounding boxes with shape (N, 4) where N is number of masks after NMS.
+        shape (tuple): Input image size as (height, width).
+        upsample (bool): Whether to upsample masks to original image size.
+
+    Returns:
+        (torch.Tensor): A binary mask tensor of shape [n, h, w], where n is the number of masks after NMS, and h and w
+            are the height and width of the input image. The mask is applied to the bounding boxes.
+    """
+    c, mh, mw = protos.shape  # CHW
+    masks = (masks_in @ protos.float().view(c, -1)).view(-1, mh, mw)  # NHW
+
+    width_ratio = mw / shape[1]
+    height_ratio = mh / shape[0]
+    ratios = torch.tensor([[width_ratio, height_ratio, width_ratio, height_ratio]], device=bboxes.device)
+
+    masks = crop_mask(masks, boxes=bboxes * ratios)  # NHW
+    if upsample:
+        masks = F.interpolate(masks[None], shape, mode="bilinear")[0]  # NHW
+    return masks.gt_(0.0).byte()
+
+
+def process_mask_native(protos, masks_in, bboxes, shape):
+    """Apply masks to bounding boxes using mask head output with native upsampling.
+
+    Args:
+        protos (torch.Tensor): Mask prototypes with shape (mask_dim, mask_h, mask_w).
+        masks_in (torch.Tensor): Mask coefficients with shape (N, mask_dim) where N is number of masks after NMS.
+        bboxes (torch.Tensor): Bounding boxes with shape (N, 4) where N is number of masks after NMS.
+        shape (tuple): Input image size as (height, width).
+
+    Returns:
+        (torch.Tensor): Binary mask tensor with shape (N, H, W).
+    """
+    c, mh, mw = protos.shape  # CHW
+    masks = (masks_in @ protos.float().view(c, -1)).view(-1, mh, mw)
+    masks = scale_masks(masks[None], shape)[0]  # NHW
+    masks = crop_mask(masks, bboxes)  # NHW
+    return masks.gt_(0.0).byte()
+
+
+def scale_masks(
+    masks: torch.Tensor,
+    shape: tuple[int, int],
+    ratio_pad: tuple[tuple[int, int], tuple[int, int]] | None = None,
+    padding: bool = True,
+) -> torch.Tensor:
+    """Rescale segment masks to target shape.
+
+    Args:
+        masks (torch.Tensor): Masks with shape (N, C, H, W).
+        shape (tuple[int, int]): Target height and width as (height, width).
+        ratio_pad (tuple, optional): Ratio and padding values as ((ratio_h, ratio_w), (pad_w, pad_h)).
+        padding (bool): Whether masks are based on YOLO-style augmented images with padding.
+
+    Returns:
+        (torch.Tensor): Rescaled masks.
+    """
+    im1_h, im1_w = masks.shape[2:]
+    im0_h, im0_w = shape[:2]
+    if im1_h == im0_h and im1_w == im0_w:
+        return masks
+
+    if ratio_pad is None:  # calculate from im0_shape
+        gain = min(im1_h / im0_h, im1_w / im0_w)  # gain  = old / new
+        pad_w, pad_h = (im1_w - im0_w * gain), (im1_h - im0_h * gain)  # wh padding
+        if padding:
+            pad_w /= 2
+            pad_h /= 2
+    else:
+        pad_w, pad_h = ratio_pad[1]
+    top, left = (round(pad_h - 0.1), round(pad_w - 0.1)) if padding else (0, 0)
+    bottom = im1_h - round(pad_h + 0.1)
+    right = im1_w - round(pad_w + 0.1)
+    return F.interpolate(masks[..., top:bottom, left:right].float(), shape, mode="bilinear")  # NCHW masks
+
+
+def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None, normalize: bool = False, padding: bool = True):
+    """Rescale segment coordinates from img1_shape to img0_shape.
+
+    Args:
+        img1_shape (tuple): Source image shape as HWC or HW (supports both).
+        coords (torch.Tensor): Coordinates to scale with shape (N, 2).
+        img0_shape (tuple): Image 0 shape as HWC or HW (supports both).
+        ratio_pad (tuple, optional): Ratio and padding values as ((ratio_h, ratio_w), (pad_w, pad_h)).
+        normalize (bool): Whether to normalize coordinates to range [0, 1].
+        padding (bool): Whether coordinates are based on YOLO-style augmented images with padding.
+
+    Returns:
+        (torch.Tensor): Scaled coordinates.
+    """
+    img0_h, img0_w = img0_shape[:2]  # supports both HWC or HW shapes
+    if ratio_pad is None:  # calculate from img0_shape
+        img1_h, img1_w = img1_shape[:2]  # supports both HWC or HW shapes
+        gain = min(img1_h / img0_h, img1_w / img0_w)  # gain  = old / new
+        pad = (img1_w - img0_w * gain) / 2, (img1_h - img0_h * gain) / 2  # wh padding
+    else:
+        gain = ratio_pad[0][0]
+        pad = ratio_pad[1]
+
+    if padding:
+        coords[..., 0] -= pad[0]  # x padding
+        coords[..., 1] -= pad[1]  # y padding
+    coords[..., 0] /= gain
+    coords[..., 1] /= gain
+    coords = clip_coords(coords, img0_shape)
+    if normalize:
+        coords[..., 0] /= img0_w  # width
+        coords[..., 1] /= img0_h  # height
+    return coords
+
+
+def regularize_rboxes(rboxes):
+    """Regularize rotated bounding boxes to range [0, pi/2).
+
+    Args:
+        rboxes (torch.Tensor): Input rotated boxes with shape (N, 5) in xywhr format.
+
+    Returns:
+        (torch.Tensor): Regularized rotated boxes.
+    """
+    x, y, w, h, t = rboxes.unbind(dim=-1)
+    # Swap edge if t >= pi/2 while not being symmetrically opposite
+    swap = t % math.pi >= math.pi / 2
+    w_ = torch.where(swap, h, w)
+    h_ = torch.where(swap, w, h)
+    t = t % (math.pi / 2)
+    return torch.stack([x, y, w_, h_, t], dim=-1)  # regularized boxes
+
+
+def masks2segments(masks: np.ndarray | torch.Tensor, strategy: str = "all") -> list[np.ndarray]:
+    """Convert masks to segments using contour detection.
+
+    Args:
+        masks (np.ndarray | torch.Tensor): Binary masks with shape (N, H, W).
+        strategy (str): Segmentation strategy, either 'all' or 'largest'.
+
+    Returns:
+        (list): List of segment masks as float32 arrays.
+    """
+    from ultralytics.data.converter import merge_multi_segment
+
+    masks = masks.astype("uint8") if isinstance(masks, np.ndarray) else masks.byte().cpu().numpy()
+    segments = []
+    for x in np.ascontiguousarray(masks):
+        c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
+        if c:
+            if strategy == "all":  # merge and concatenate all segments
+                c = (
+                    np.concatenate(merge_multi_segment([x.reshape(-1, 2) for x in c]))
+                    if len(c) > 1
+                    else c[0].reshape(-1, 2)
+                )
+            elif strategy == "largest":  # select largest segment
+                c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
+        else:
+            c = np.zeros((0, 2))  # no segments found
+        segments.append(c.astype("float32"))
+    return segments
+
+
+def convert_torch2numpy_batch(batch: torch.Tensor) -> np.ndarray:
+    """Convert a batch of FP32 torch tensors to NumPy uint8 arrays, changing from BCHW to BHWC layout.
+
+    Args:
+        batch (torch.Tensor): Input tensor batch with shape (Batch, Channels, Height, Width) and dtype torch.float32.
+
+    Returns:
+        (np.ndarray): Output NumPy array batch with shape (Batch, Height, Width, Channels) and dtype uint8.
+    """
+    return (batch.permute(0, 2, 3, 1).contiguous() * 255).clamp(0, 255).byte().cpu().numpy()
+
+
+def clean_str(s):
+    """Clean a string by replacing special characters with '_' character.
+
+    Args:
+        s (str): A string needing special characters replaced.
+
+    Returns:
+        (str): A string with special characters replaced by an underscore _.
+    """
+    return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨`><+]", repl="_", string=s)
+
+
+def empty_like(x):
+    """Create empty torch.Tensor or np.ndarray with same shape and dtype as input."""
+    return torch.empty_like(x, dtype=x.dtype) if isinstance(x, torch.Tensor) else np.empty_like(x, dtype=x.dtype)

ultralytics/utils/patches.py

@@ -0,0 +1,266 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""Monkey patches to update/extend functionality of existing functions."""
+
+from __future__ import annotations
+
+import time
+from contextlib import contextmanager
+from copy import copy
+from pathlib import Path
+from typing import Any
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+
+# OpenCV Multilanguage-friendly functions ------------------------------------------------------------------------------
+_imshow = cv2.imshow  # copy to avoid recursion errors
+
+
+def imread(filename: str, flags: int = cv2.IMREAD_COLOR) -> np.ndarray | None:
+    """Read an image from a file with multilanguage filename support.
+
+    Args:
+        filename (str): Path to the file to read.
+        flags (int, optional): Flag that can take values of cv2.IMREAD_*. Controls how the image is read.
+
+    Returns:
+        (np.ndarray | None): The read image array, or None if reading fails.
+
+    Examples:
+        >>> img = imread("path/to/image.jpg")
+        >>> img = imread("path/to/image.jpg", cv2.IMREAD_GRAYSCALE)
+    """
+    file_bytes = np.fromfile(filename, np.uint8)
+    if filename.endswith((".tiff", ".tif")):
+        success, frames = cv2.imdecodemulti(file_bytes, cv2.IMREAD_UNCHANGED)
+        if success:
+            # Handle multi-frame TIFFs and color images
+            return frames[0] if len(frames) == 1 and frames[0].ndim == 3 else np.stack(frames, axis=2)
+        return None
+    else:
+        im = cv2.imdecode(file_bytes, flags)
+        # Fallback for formats OpenCV imdecode may not support (AVIF, HEIC)
+        if im is None and filename.lower().endswith((".avif", ".heic")):
+            im = _imread_pil(filename, flags)
+        return im[..., None] if im is not None and im.ndim == 2 else im  # Always ensure 3 dimensions
+
+
+# PIL patches ---------------------------------------------------------------------------------------------------------
+_image_open = Image.open  # copy to avoid recursion errors
+_pil_plugins_registered = False
+
+
+def image_open(filename, *args, **kwargs):
+    """Open an image with PIL, lazily registering the HEIF plugin on first failure.
+
+    This monkey-patches PIL.Image.open to add HEIC/HEIF support via pi-heif (lightweight, decode-only), avoiding the
+    ~800ms startup cost of importing the package unless actually needed. AVIF is supported natively by Pillow 12+ and
+    does not require a plugin.
+
+    Args:
+        filename (str): Path to the image file.
+        *args (Any): Additional positional arguments passed to PIL.Image.open.
+        **kwargs (Any): Additional keyword arguments passed to PIL.Image.open.
+
+    Returns:
+        (PIL.Image.Image): The opened PIL image.
+    """
+    global _pil_plugins_registered
+    if _pil_plugins_registered:
+        return _image_open(filename, *args, **kwargs)
+    try:
+        return _image_open(filename, *args, **kwargs)
+    except Exception as e:
+        suffix = Path(filename).suffix.lower() if isinstance(filename, (str, Path)) else ""
+        if suffix not in {".heic", ".heif", ".hif"}:
+            raise e
+        from ultralytics.utils.checks import check_requirements
+
+        check_requirements("pi-heif")
+        from pi_heif import register_heif_opener
+
+        register_heif_opener()
+        _pil_plugins_registered = True
+        return _image_open(filename, *args, **kwargs)
+
+
+Image.open = image_open  # apply patch
+
+
+def _imread_pil(filename: str, flags: int = cv2.IMREAD_COLOR) -> np.ndarray | None:
+    """Read an image using PIL as fallback for formats not supported by OpenCV.
+
+    Args:
+        filename (str): Path to the file to read.
+        flags (int, optional): OpenCV imread flags (used to determine grayscale conversion).
+
+    Returns:
+        (np.ndarray | None): The read image array in BGR format, or None if reading fails.
+    """
+    try:
+        with Image.open(filename) as img:
+            if flags == cv2.IMREAD_GRAYSCALE:
+                return np.asarray(img.convert("L"))
+            return cv2.cvtColor(np.asarray(img.convert("RGB")), cv2.COLOR_RGB2BGR)
+    except Exception:
+        return None
+
+
+def imwrite(filename: str, img: np.ndarray, params: list[int] | None = None) -> bool:
+    """Write an image to a file with multilanguage filename support.
+
+    Args:
+        filename (str): Path to the file to write.
+        img (np.ndarray): Image to write.
+        params (list[int], optional): Additional parameters for image encoding.
+
+    Returns:
+        (bool): True if the file was written successfully, False otherwise.
+
+    Examples:
+        >>> import numpy as np
+        >>> img = np.zeros((100, 100, 3), dtype=np.uint8)  # Create a black image
+        >>> success = imwrite("output.jpg", img)  # Write image to file
+        >>> print(success)
+        True
+    """
+    try:
+        cv2.imencode(Path(filename).suffix, img, params)[1].tofile(filename)
+        return True
+    except Exception:
+        return False
+
+
+def imshow(winname: str, mat: np.ndarray) -> None:
+    """Display an image in the specified window with multilanguage window name support.
+
+    This function is a wrapper around OpenCV's imshow function that displays an image in a named window. It handles
+    multilanguage window names by encoding them properly for OpenCV compatibility.
+
+    Args:
+        winname (str): Name of the window where the image will be displayed. If a window with this name already exists,
+            the image will be displayed in that window.
+        mat (np.ndarray): Image to be shown. Should be a valid numpy array representing an image.
+
+    Examples:
+        >>> import numpy as np
+        >>> img = np.zeros((300, 300, 3), dtype=np.uint8)  # Create a black image
+        >>> img[:100, :100] = [255, 0, 0]  # Add a blue square
+        >>> imshow("Example Window", img)  # Display the image
+    """
+    _imshow(winname.encode("unicode_escape").decode(), mat)
+
+
+# PyTorch functions ----------------------------------------------------------------------------------------------------
+_torch_save = torch.save
+
+
+def torch_load(*args, **kwargs):
+    """Load a PyTorch model with updated arguments to avoid warnings.
+
+    This function wraps torch.load and adds the 'weights_only' argument for PyTorch 1.13.0+ to prevent warnings.
+
+    Args:
+        *args (Any): Variable length argument list to pass to torch.load.
+        **kwargs (Any): Arbitrary keyword arguments to pass to torch.load.
+
+    Returns:
+        (Any): The loaded PyTorch object.
+
+    Notes:
+        For PyTorch versions 1.13 and above, this function automatically sets `weights_only=False` if the argument is
+        not provided, to avoid deprecation warnings.
+    """
+    from ultralytics.utils.torch_utils import TORCH_1_13
+
+    if TORCH_1_13 and "weights_only" not in kwargs:
+        kwargs["weights_only"] = False
+
+    return torch.load(*args, **kwargs)
+
+
+def torch_save(*args, **kwargs):
+    """Save PyTorch objects with retry mechanism for robustness.
+
+    This function wraps torch.save with 3 retries and exponential backoff in case of save failures, which can occur due
+    to device flushing delays or antivirus scanning.
+
+    Args:
+        *args (Any): Positional arguments to pass to torch.save.
+        **kwargs (Any): Keyword arguments to pass to torch.save.
+
+    Examples:
+        >>> model = torch.nn.Linear(10, 1)
+        >>> torch_save(model.state_dict(), "model.pt")
+    """
+    for i in range(4):  # 3 retries
+        try:
+            return _torch_save(*args, **kwargs)
+        except RuntimeError as e:  # Unable to save, possibly waiting for device to flush or antivirus scan
+            if i == 3:
+                raise e
+            time.sleep((2**i) / 2)  # Exponential backoff: 0.5s, 1.0s, 2.0s
+
+
+@contextmanager
+def arange_patch(args):
+    """Workaround for ONNX torch.arange incompatibility with FP16.
+
+    https://github.com/pytorch/pytorch/issues/148041.
+    """
+    if args.dynamic and args.half and args.format == "onnx":
+        func = torch.arange
+
+        def arange(*args, dtype=None, **kwargs):
+            """Wrap torch.arange to cast dtype after creation instead of passing it directly."""
+            return func(*args, **kwargs).to(dtype)  # cast to dtype instead of passing dtype
+
+        torch.arange = arange  # patch
+        yield
+        torch.arange = func  # unpatch
+    else:
+        yield
+
+
+@contextmanager
+def onnx_export_patch():
+    """Workaround for ONNX export issues in PyTorch 2.9+ with Dynamo enabled."""
+    from ultralytics.utils.torch_utils import TORCH_2_9
+
+    if TORCH_2_9:
+        func = torch.onnx.export
+
+        def torch_export(*args, **kwargs):
+            """Export model to ONNX format with Dynamo disabled for compatibility."""
+            return func(*args, **kwargs, dynamo=False)
+
+        torch.onnx.export = torch_export  # patch
+        yield
+        torch.onnx.export = func  # unpatch
+    else:
+        yield
+
+
+@contextmanager
+def override_configs(args, overrides: dict[str, Any] | None = None):
+    """Context manager to temporarily override configurations in args.
+
+    Args:
+        args (IterableSimpleNamespace): Original configuration arguments.
+        overrides (dict[str, Any] | None): Dictionary of overrides to apply.
+
+    Yields:
+        (IterableSimpleNamespace): Configuration arguments with overrides applied.
+    """
+    if overrides:
+        original_args = copy(args)
+        for key, value in overrides.items():
+            setattr(args, key, value)
+        try:
+            yield args
+        finally:
+            args.__dict__.update(original_args.__dict__)
+    else:
+        yield args

ultralytics/utils/tal.py

@@ -0,0 +1,493 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from . import LOGGER
+from .metrics import bbox_iou, probiou
+from .ops import xywh2xyxy, xywhr2xyxyxyxy, xyxy2xywh
+from .torch_utils import TORCH_1_11
+
+
+class TaskAlignedAssigner(nn.Module):
+    """A task-aligned assigner for object detection.
+
+    This class assigns ground-truth (gt) objects to anchors based on the task-aligned metric, which combines both
+    classification and localization information.
+
+    Attributes:
+        topk (int): The number of top candidates to consider.
+        topk2 (int): Secondary topk value for additional filtering.
+        num_classes (int): The number of object classes.
+        alpha (float): The alpha parameter for the classification component of the task-aligned metric.
+        beta (float): The beta parameter for the localization component of the task-aligned metric.
+        stride (list): List of stride values for different feature levels.
+        stride_val (int): The stride value used for select_candidates_in_gts.
+        eps (float): A small value to prevent division by zero.
+    """
+
+    def __init__(
+        self,
+        topk: int = 13,
+        num_classes: int = 80,
+        alpha: float = 1.0,
+        beta: float = 6.0,
+        stride: list = [8, 16, 32],
+        eps: float = 1e-9,
+        topk2=None,
+    ):
+        """Initialize a TaskAlignedAssigner object with customizable hyperparameters.
+
+        Args:
+            topk (int, optional): The number of top candidates to consider.
+            num_classes (int, optional): The number of object classes.
+            alpha (float, optional): The alpha parameter for the classification component of the task-aligned metric.
+            beta (float, optional): The beta parameter for the localization component of the task-aligned metric.
+            stride (list, optional): List of stride values for different feature levels.
+            eps (float, optional): A small value to prevent division by zero.
+            topk2 (int, optional): Secondary topk value for additional filtering.
+        """
+        super().__init__()
+        self.topk = topk
+        self.topk2 = topk2 or topk
+        self.num_classes = num_classes
+        self.alpha = alpha
+        self.beta = beta
+        self.stride = stride
+        self.stride_val = self.stride[1] if len(self.stride) > 1 else self.stride[0]
+        self.eps = eps
+
+    @torch.no_grad()
+    def forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt):
+        """Compute the task-aligned assignment.
+
+        Args:
+            pd_scores (torch.Tensor): Predicted classification scores with shape (bs, num_total_anchors, num_classes).
+            pd_bboxes (torch.Tensor): Predicted bounding boxes with shape (bs, num_total_anchors, 4).
+            anc_points (torch.Tensor): Anchor points with shape (num_total_anchors, 2).
+            gt_labels (torch.Tensor): Ground truth labels with shape (bs, n_max_boxes, 1).
+            gt_bboxes (torch.Tensor): Ground truth boxes with shape (bs, n_max_boxes, 4).
+            mask_gt (torch.Tensor): Mask for valid ground truth boxes with shape (bs, n_max_boxes, 1).
+
+        Returns:
+            target_labels (torch.Tensor): Target labels with shape (bs, num_total_anchors).
+            target_bboxes (torch.Tensor): Target bounding boxes with shape (bs, num_total_anchors, 4).
+            target_scores (torch.Tensor): Target scores with shape (bs, num_total_anchors, num_classes).
+            fg_mask (torch.Tensor): Foreground mask with shape (bs, num_total_anchors).
+            target_gt_idx (torch.Tensor): Target ground truth indices with shape (bs, num_total_anchors).
+
+        References:
+            https://github.com/Nioolek/PPYOLOE_pytorch/blob/master/ppyoloe/assigner/tal_assigner.py
+        """
+        self.bs = pd_scores.shape[0]
+        self.n_max_boxes = gt_bboxes.shape[1]
+        device = gt_bboxes.device
+
+        if self.n_max_boxes == 0:
+            return (
+                torch.full_like(pd_scores[..., 0], self.num_classes),
+                torch.zeros_like(pd_bboxes),
+                torch.zeros_like(pd_scores),
+                torch.zeros_like(pd_scores[..., 0]),
+                torch.zeros_like(pd_scores[..., 0]),
+            )
+
+        try:
+            return self._forward(pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt)
+        except RuntimeError as e:
+            if "out of memory" in str(e).lower():
+                # Move tensors to CPU, compute, then move back to original device
+                LOGGER.warning("CUDA OutOfMemoryError in TaskAlignedAssigner, using CPU")
+                cpu_tensors = [t.cpu() for t in (pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt)]
+                result = self._forward(*cpu_tensors)
+                return tuple(t.to(device) for t in result)
+            raise
+
+    def _forward(self, pd_scores, pd_bboxes, anc_points, gt_labels, gt_bboxes, mask_gt):
+        """Compute the task-aligned assignment.
+
+        Args:
+            pd_scores (torch.Tensor): Predicted classification scores with shape (bs, num_total_anchors, num_classes).
+            pd_bboxes (torch.Tensor): Predicted bounding boxes with shape (bs, num_total_anchors, 4).
+            anc_points (torch.Tensor): Anchor points with shape (num_total_anchors, 2).
+            gt_labels (torch.Tensor): Ground truth labels with shape (bs, n_max_boxes, 1).
+            gt_bboxes (torch.Tensor): Ground truth boxes with shape (bs, n_max_boxes, 4).
+            mask_gt (torch.Tensor): Mask for valid ground truth boxes with shape (bs, n_max_boxes, 1).
+
+        Returns:
+            target_labels (torch.Tensor): Target labels with shape (bs, num_total_anchors).
+            target_bboxes (torch.Tensor): Target bounding boxes with shape (bs, num_total_anchors, 4).
+            target_scores (torch.Tensor): Target scores with shape (bs, num_total_anchors, num_classes).
+            fg_mask (torch.Tensor): Foreground mask with shape (bs, num_total_anchors).
+            target_gt_idx (torch.Tensor): Target ground truth indices with shape (bs, num_total_anchors).
+        """
+        mask_pos, align_metric, overlaps = self.get_pos_mask(
+            pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, mask_gt
+        )
+
+        target_gt_idx, fg_mask, mask_pos = self.select_highest_overlaps(
+            mask_pos, overlaps, self.n_max_boxes, align_metric
+        )
+
+        # Assigned target
+        target_labels, target_bboxes, target_scores = self.get_targets(gt_labels, gt_bboxes, target_gt_idx, fg_mask)
+
+        # Normalize
+        align_metric *= mask_pos
+        pos_align_metrics = align_metric.amax(dim=-1, keepdim=True)  # b, max_num_obj
+        pos_overlaps = (overlaps * mask_pos).amax(dim=-1, keepdim=True)  # b, max_num_obj
+        norm_align_metric = (align_metric * pos_overlaps / (pos_align_metrics + self.eps)).amax(-2).unsqueeze(-1)
+        target_scores = target_scores * norm_align_metric
+
+        return target_labels, target_bboxes, target_scores, fg_mask.bool(), target_gt_idx
+
+    def get_pos_mask(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, mask_gt):
+        """Get positive mask for each ground truth box.
+
+        Args:
+            pd_scores (torch.Tensor): Predicted classification scores with shape (bs, num_total_anchors, num_classes).
+            pd_bboxes (torch.Tensor): Predicted bounding boxes with shape (bs, num_total_anchors, 4).
+            gt_labels (torch.Tensor): Ground truth labels with shape (bs, n_max_boxes, 1).
+            gt_bboxes (torch.Tensor): Ground truth boxes with shape (bs, n_max_boxes, 4).
+            anc_points (torch.Tensor): Anchor points with shape (num_total_anchors, 2).
+            mask_gt (torch.Tensor): Mask for valid ground truth boxes with shape (bs, n_max_boxes, 1).
+
+        Returns:
+            mask_pos (torch.Tensor): Positive mask with shape (bs, max_num_obj, h*w).
+            align_metric (torch.Tensor): Alignment metric with shape (bs, max_num_obj, h*w).
+            overlaps (torch.Tensor): Overlaps between predicted vs ground truth boxes with shape (bs, max_num_obj, h*w).
+        """
+        mask_in_gts = self.select_candidates_in_gts(anc_points, gt_bboxes, mask_gt)
+        # Get anchor_align metric, (b, max_num_obj, h*w)
+        align_metric, overlaps = self.get_box_metrics(pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_in_gts * mask_gt)
+        # Get topk_metric mask, (b, max_num_obj, h*w)
+        mask_topk = self.select_topk_candidates(align_metric, topk_mask=mask_gt.expand(-1, -1, self.topk).bool())
+        # Merge all mask to a final mask, (b, max_num_obj, h*w)
+        mask_pos = mask_topk * mask_in_gts * mask_gt
+
+        return mask_pos, align_metric, overlaps
+
+    def get_box_metrics(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, mask_gt):
+        """Compute alignment metric given predicted and ground truth bounding boxes.
+
+        Args:
+            pd_scores (torch.Tensor): Predicted classification scores with shape (bs, num_total_anchors, num_classes).
+            pd_bboxes (torch.Tensor): Predicted bounding boxes with shape (bs, num_total_anchors, 4).
+            gt_labels (torch.Tensor): Ground truth labels with shape (bs, n_max_boxes, 1).
+            gt_bboxes (torch.Tensor): Ground truth boxes with shape (bs, n_max_boxes, 4).
+            mask_gt (torch.Tensor): Mask for valid ground truth boxes with shape (bs, n_max_boxes, h*w).
+
+        Returns:
+            align_metric (torch.Tensor): Alignment metric combining classification and localization.
+            overlaps (torch.Tensor): IoU overlaps between predicted and ground truth boxes.
+        """
+        na = pd_bboxes.shape[-2]
+        mask_gt = mask_gt.bool()  # b, max_num_obj, h*w
+        overlaps = torch.zeros([self.bs, self.n_max_boxes, na], dtype=pd_bboxes.dtype, device=pd_bboxes.device)
+        bbox_scores = torch.zeros([self.bs, self.n_max_boxes, na], dtype=pd_scores.dtype, device=pd_scores.device)
+
+        ind = torch.zeros([2, self.bs, self.n_max_boxes], dtype=torch.long)  # 2, b, max_num_obj
+        ind[0] = torch.arange(end=self.bs).view(-1, 1).expand(-1, self.n_max_boxes)  # b, max_num_obj
+        ind[1] = gt_labels.squeeze(-1)  # b, max_num_obj
+        # Get the scores of each grid for each gt cls
+        bbox_scores[mask_gt] = pd_scores[ind[0], :, ind[1]][mask_gt]  # b, max_num_obj, h*w
+
+        # (b, max_num_obj, 1, 4), (b, 1, h*w, 4)
+        pd_boxes = pd_bboxes.unsqueeze(1).expand(-1, self.n_max_boxes, -1, -1)[mask_gt]
+        gt_boxes = gt_bboxes.unsqueeze(2).expand(-1, -1, na, -1)[mask_gt]
+        overlaps[mask_gt] = self.iou_calculation(gt_boxes, pd_boxes)
+
+        align_metric = bbox_scores.pow(self.alpha) * overlaps.pow(self.beta)
+        return align_metric, overlaps
+
+    def iou_calculation(self, gt_bboxes, pd_bboxes):
+        """Calculate IoU for horizontal bounding boxes.
+
+        Args:
+            gt_bboxes (torch.Tensor): Ground truth boxes.
+            pd_bboxes (torch.Tensor): Predicted boxes.
+
+        Returns:
+            (torch.Tensor): IoU values between each pair of boxes.
+        """
+        return bbox_iou(gt_bboxes, pd_bboxes, xywh=False, CIoU=True).squeeze(-1).clamp_(0)
+
+    def select_topk_candidates(self, metrics, topk_mask=None):
+        """Select the top-k candidates based on the given metrics.
+
+        Args:
+            metrics (torch.Tensor): A tensor of shape (b, max_num_obj, h*w), where b is the batch size, max_num_obj is
+                the maximum number of objects, and h*w represents the total number of anchor points.
+            topk_mask (torch.Tensor, optional): An optional boolean tensor of shape (b, max_num_obj, topk), where topk
+                is the number of top candidates to consider. If not provided, the top-k values are automatically
+                computed based on the given metrics.
+
+        Returns:
+            (torch.Tensor): A tensor of shape (b, max_num_obj, h*w) containing the selected top-k candidates.
+        """
+        # (b, max_num_obj, topk)
+        topk_metrics, topk_idxs = torch.topk(metrics, self.topk, dim=-1, largest=True)
+        if topk_mask is None:
+            topk_mask = (topk_metrics.max(-1, keepdim=True)[0] > self.eps).expand_as(topk_idxs)
+        # (b, max_num_obj, topk)
+        topk_idxs.masked_fill_(~topk_mask, 0)
+
+        # (b, max_num_obj, topk, h*w) -> (b, max_num_obj, h*w)
+        count_tensor = torch.zeros(metrics.shape, dtype=torch.int8, device=topk_idxs.device)
+        ones = torch.ones_like(topk_idxs[:, :, :1], dtype=torch.int8, device=topk_idxs.device)
+        for k in range(self.topk):
+            # Expand topk_idxs for each value of k and add 1 at the specified positions
+            count_tensor.scatter_add_(-1, topk_idxs[:, :, k : k + 1], ones)
+        # Filter invalid bboxes
+        count_tensor.masked_fill_(count_tensor > 1, 0)
+
+        return count_tensor.to(metrics.dtype)
+
+    def get_targets(self, gt_labels, gt_bboxes, target_gt_idx, fg_mask):
+        """Compute target labels, target bounding boxes, and target scores for the positive anchor points.
+
+        Args:
+            gt_labels (torch.Tensor): Ground truth labels of shape (b, max_num_obj, 1), where b is the batch size and
+                max_num_obj is the maximum number of objects.
+            gt_bboxes (torch.Tensor): Ground truth bounding boxes of shape (b, max_num_obj, 4).
+            target_gt_idx (torch.Tensor): Indices of the assigned ground truth objects for positive anchor points, with
+                shape (b, h*w), where h*w is the total number of anchor points.
+            fg_mask (torch.Tensor): A boolean tensor of shape (b, h*w) indicating the positive (foreground) anchor
+                points.
+
+        Returns:
+            target_labels (torch.Tensor): Target labels for positive anchor points with shape (b, h*w).
+            target_bboxes (torch.Tensor): Target bounding boxes for positive anchor points with shape (b, h*w, 4).
+            target_scores (torch.Tensor): Target scores for positive anchor points with shape (b, h*w, num_classes).
+        """
+        # Assigned target labels, (b, 1)
+        batch_ind = torch.arange(end=self.bs, dtype=torch.int64, device=gt_labels.device)[..., None]
+        target_gt_idx = target_gt_idx + batch_ind * self.n_max_boxes  # (b, h*w)
+        target_labels = gt_labels.long().flatten()[target_gt_idx]  # (b, h*w)
+
+        # Assigned target boxes, (b, max_num_obj, 4) -> (b, h*w, 4)
+        target_bboxes = gt_bboxes.view(-1, gt_bboxes.shape[-1])[target_gt_idx]
+
+        # Assigned target scores
+        target_labels.clamp_(0)
+
+        # 10x faster than F.one_hot()
+        target_scores = torch.zeros(
+            (target_labels.shape[0], target_labels.shape[1], self.num_classes),
+            dtype=torch.int64,
+            device=target_labels.device,
+        )  # (b, h*w, 80)
+        target_scores.scatter_(2, target_labels.unsqueeze(-1), 1)
+
+        fg_scores_mask = fg_mask[:, :, None].repeat(1, 1, self.num_classes)  # (b, h*w, 80)
+        target_scores = torch.where(fg_scores_mask > 0, target_scores, 0)
+
+        return target_labels, target_bboxes, target_scores
+
+    def select_candidates_in_gts(self, xy_centers, gt_bboxes, mask_gt, eps=1e-9):
+        """Select positive anchor centers within ground truth bounding boxes.
+
+        Args:
+            xy_centers (torch.Tensor): Anchor center coordinates, shape (h*w, 2).
+            gt_bboxes (torch.Tensor): Ground truth bounding boxes, shape (b, n_boxes, 4).
+            mask_gt (torch.Tensor): Mask for valid ground truth boxes, shape (b, n_boxes, 1).
+            eps (float, optional): Small value for numerical stability.
+
+        Returns:
+            (torch.Tensor): Boolean mask of positive anchors, shape (b, n_boxes, h*w).
+
+        Notes:
+            - b: batch size, n_boxes: number of ground truth boxes, h: height, w: width.
+            - Bounding box format: [x_min, y_min, x_max, y_max].
+        """
+        gt_bboxes_xywh = xyxy2xywh(gt_bboxes)
+        wh_mask = gt_bboxes_xywh[..., 2:] < self.stride[0]  # the smallest stride
+        gt_bboxes_xywh[..., 2:] = torch.where(
+            (wh_mask * mask_gt).bool(),
+            torch.tensor(self.stride_val, dtype=gt_bboxes_xywh.dtype, device=gt_bboxes_xywh.device),
+            gt_bboxes_xywh[..., 2:],
+        )
+        gt_bboxes = xywh2xyxy(gt_bboxes_xywh)
+
+        n_anchors = xy_centers.shape[0]
+        bs, n_boxes, _ = gt_bboxes.shape
+        lt, rb = gt_bboxes.view(-1, 1, 4).chunk(2, 2)  # left-top, right-bottom
+        bbox_deltas = torch.cat((xy_centers[None] - lt, rb - xy_centers[None]), dim=2).view(bs, n_boxes, n_anchors, -1)
+        return bbox_deltas.amin(3).gt_(eps)
+
+    def select_highest_overlaps(self, mask_pos, overlaps, n_max_boxes, align_metric):
+        """Select anchor boxes with highest IoU when assigned to multiple ground truths.
+
+        Args:
+            mask_pos (torch.Tensor): Positive mask, shape (b, n_max_boxes, h*w).
+            overlaps (torch.Tensor): IoU overlaps, shape (b, n_max_boxes, h*w).
+            n_max_boxes (int): Maximum number of ground truth boxes.
+            align_metric (torch.Tensor): Alignment metric for selecting best matches.
+
+        Returns:
+            target_gt_idx (torch.Tensor): Indices of assigned ground truths, shape (b, h*w).
+            fg_mask (torch.Tensor): Foreground mask, shape (b, h*w).
+            mask_pos (torch.Tensor): Updated positive mask, shape (b, n_max_boxes, h*w).
+        """
+        # Convert (b, n_max_boxes, h*w) -> (b, h*w)
+        fg_mask = mask_pos.sum(-2)
+        if fg_mask.max() > 1:  # one anchor is assigned to multiple gt_bboxes
+            mask_multi_gts = (fg_mask.unsqueeze(1) > 1).expand(-1, n_max_boxes, -1)  # (b, n_max_boxes, h*w)
+
+            max_overlaps_idx = overlaps.argmax(1)  # (b, h*w)
+            is_max_overlaps = torch.zeros(mask_pos.shape, dtype=mask_pos.dtype, device=mask_pos.device)
+            is_max_overlaps.scatter_(1, max_overlaps_idx.unsqueeze(1), 1)
+            mask_pos = torch.where(mask_multi_gts, is_max_overlaps, mask_pos).float()  # (b, n_max_boxes, h*w)
+
+            fg_mask = mask_pos.sum(-2)
+
+        if self.topk2 != self.topk:
+            align_metric = align_metric * mask_pos  # update overlaps
+            max_overlaps_idx = torch.topk(align_metric, self.topk2, dim=-1, largest=True).indices  # (b, n_max_boxes)
+            topk_idx = torch.zeros(mask_pos.shape, dtype=mask_pos.dtype, device=mask_pos.device)  # update mask_pos
+            topk_idx.scatter_(-1, max_overlaps_idx, 1.0)
+            mask_pos *= topk_idx
+            fg_mask = mask_pos.sum(-2)
+        # Find each grid serve which gt(index)
+        target_gt_idx = mask_pos.argmax(-2)  # (b, h*w)
+        return target_gt_idx, fg_mask, mask_pos
+
+
+class RotatedTaskAlignedAssigner(TaskAlignedAssigner):
+    """Assigns ground-truth objects to rotated bounding boxes using a task-aligned metric."""
+
+    def iou_calculation(self, gt_bboxes, pd_bboxes):
+        """Calculate IoU for rotated bounding boxes."""
+        return probiou(gt_bboxes, pd_bboxes).squeeze(-1).clamp_(0)
+
+    def select_candidates_in_gts(self, xy_centers, gt_bboxes, mask_gt):
+        """Select the positive anchor center in gt for rotated bounding boxes.
+
+        Args:
+            xy_centers (torch.Tensor): Anchor center coordinates with shape (h*w, 2).
+            gt_bboxes (torch.Tensor): Ground truth bounding boxes with shape (b, n_boxes, 5).
+            mask_gt (torch.Tensor): Mask for valid ground truth boxes with shape (b, n_boxes, 1).
+
+        Returns:
+            (torch.Tensor): Boolean mask of positive anchors with shape (b, n_boxes, h*w).
+        """
+        wh_mask = gt_bboxes[..., 2:4] < self.stride[0]
+        gt_bboxes[..., 2:4] = torch.where(
+            (wh_mask * mask_gt).bool(),
+            torch.tensor(self.stride_val, dtype=gt_bboxes.dtype, device=gt_bboxes.device),
+            gt_bboxes[..., 2:4],
+        )
+
+        # (b, n_boxes, 5) --> (b, n_boxes, 4, 2)
+        corners = xywhr2xyxyxyxy(gt_bboxes)
+        # (b, n_boxes, 1, 2)
+        a, b, _, d = corners.split(1, dim=-2)
+        ab = b - a
+        ad = d - a
+
+        # (b, n_boxes, h*w, 2)
+        ap = xy_centers - a
+        norm_ab = (ab * ab).sum(dim=-1)
+        norm_ad = (ad * ad).sum(dim=-1)
+        ap_dot_ab = (ap * ab).sum(dim=-1)
+        ap_dot_ad = (ap * ad).sum(dim=-1)
+        return (ap_dot_ab >= 0) & (ap_dot_ab <= norm_ab) & (ap_dot_ad >= 0) & (ap_dot_ad <= norm_ad)  # is_in_box
+
+
+def make_anchors(feats, strides, grid_cell_offset=0.5):
+    """Generate anchors from features."""
+    anchor_points, stride_tensor = [], []
+    assert feats is not None
+    dtype, device = feats[0].dtype, feats[0].device
+    for i in range(len(feats)):  # use len(feats) to avoid TracerWarning from iterating over strides tensor
+        stride = strides[i]
+        h, w = feats[i].shape[2:] if isinstance(feats, list) else (int(feats[i][0]), int(feats[i][1]))
+        sx = torch.arange(end=w, device=device, dtype=dtype) + grid_cell_offset  # shift x
+        sy = torch.arange(end=h, device=device, dtype=dtype) + grid_cell_offset  # shift y
+        sy, sx = torch.meshgrid(sy, sx, indexing="ij") if TORCH_1_11 else torch.meshgrid(sy, sx)
+        anchor_points.append(torch.stack((sx, sy), -1).view(-1, 2))
+        stride_tensor.append(torch.full((h * w, 1), stride, dtype=dtype, device=device))
+    return torch.cat(anchor_points), torch.cat(stride_tensor)
+
+
+def dist2bbox(distance, anchor_points, xywh=True, dim=-1):
+    """Transform distance(ltrb) to box(xywh or xyxy)."""
+    lt, rb = distance.chunk(2, dim)
+    x1y1 = anchor_points - lt
+    x2y2 = anchor_points + rb
+    if xywh:
+        c_xy = (x1y1 + x2y2) / 2
+        wh = x2y2 - x1y1
+        return torch.cat([c_xy, wh], dim)  # xywh bbox
+    return torch.cat((x1y1, x2y2), dim)  # xyxy bbox
+
+
+def bbox2dist(anchor_points: torch.Tensor, bbox: torch.Tensor, reg_max: int | None = None) -> torch.Tensor:
+    """Transform bbox(xyxy) to dist(ltrb)."""
+    x1y1, x2y2 = bbox.chunk(2, -1)
+    dist = torch.cat((anchor_points - x1y1, x2y2 - anchor_points), -1)
+    if reg_max is not None:
+        dist = dist.clamp_(0, reg_max - 0.01)  # dist (lt, rb)
+    return dist
+
+
+def dist2rbox(pred_dist, pred_angle, anchor_points, dim=-1):
+    """Decode predicted rotated bounding box coordinates from anchor points and distribution.
+
+    Args:
+        pred_dist (torch.Tensor): Predicted rotated distance with shape (bs, h*w, 4).
+        pred_angle (torch.Tensor): Predicted angle with shape (bs, h*w, 1).
+        anchor_points (torch.Tensor): Anchor points with shape (h*w, 2).
+        dim (int, optional): Dimension along which to split.
+
+    Returns:
+        (torch.Tensor): Predicted rotated bounding boxes with shape (bs, h*w, 4).
+    """
+    lt, rb = pred_dist.split(2, dim=dim)
+    cos, sin = torch.cos(pred_angle), torch.sin(pred_angle)
+    # (bs, h*w, 1)
+    xf, yf = ((rb - lt) / 2).split(1, dim=dim)
+    x, y = xf * cos - yf * sin, xf * sin + yf * cos
+    xy = torch.cat([x, y], dim=dim) + anchor_points
+    return torch.cat([xy, lt + rb], dim=dim)
+
+
+def rbox2dist(
+    target_bboxes: torch.Tensor,
+    anchor_points: torch.Tensor,
+    target_angle: torch.Tensor,
+    dim: int = -1,
+    reg_max: int | None = None,
+):
+    """Transform rotated bounding box (xywh) to distance (ltrb). This is the inverse of dist2rbox.
+
+    Args:
+        target_bboxes (torch.Tensor): Target rotated bounding boxes with shape (bs, h*w, 4), format [x, y, w, h].
+        anchor_points (torch.Tensor): Anchor points with shape (h*w, 2).
+        target_angle (torch.Tensor): Target angle with shape (bs, h*w, 1).
+        dim (int, optional): Dimension along which to split.
+        reg_max (int, optional): Maximum regression value for clamping.
+
+    Returns:
+        (torch.Tensor): Rotated distance with shape (bs, h*w, 4), format [l, t, r, b].
+    """
+    xy, wh = target_bboxes.split(2, dim=dim)
+    offset = xy - anchor_points  # (bs, h*w, 2)
+    offset_x, offset_y = offset.split(1, dim=dim)
+    cos, sin = torch.cos(target_angle), torch.sin(target_angle)
+    xf = offset_x * cos + offset_y * sin
+    yf = -offset_x * sin + offset_y * cos
+
+    w, h = wh.split(1, dim=dim)
+    target_l = w / 2 - xf
+    target_t = h / 2 - yf
+    target_r = w / 2 + xf
+    target_b = h / 2 + yf
+
+    dist = torch.cat([target_l, target_t, target_r, target_b], dim=dim)
+    if reg_max is not None:
+        dist = dist.clamp_(0, reg_max - 0.01)
+
+    return dist

ultralytics/utils/torch_utils.py

@@ -0,0 +1,988 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import functools
+import gc
+import math
+import os
+import random
+import time
+from contextlib import contextmanager
+from copy import deepcopy
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ultralytics import __version__
+from ultralytics.utils import (
+    DEFAULT_CFG_DICT,
+    DEFAULT_CFG_KEYS,
+    LOGGER,
+    NUM_THREADS,
+    PYTHON_VERSION,
+    TORCH_VERSION,
+    TORCHVISION_VERSION,
+    WINDOWS,
+    colorstr,
+)
+from ultralytics.utils.checks import check_version
+from ultralytics.utils.cpu import CPUInfo
+from ultralytics.utils.patches import torch_load
+
+# Version checks (all default to version>=min_version)
+TORCH_1_9 = check_version(TORCH_VERSION, "1.9.0")
+TORCH_1_10 = check_version(TORCH_VERSION, "1.10.0")
+TORCH_1_11 = check_version(TORCH_VERSION, "1.11.0")
+TORCH_1_13 = check_version(TORCH_VERSION, "1.13.0")
+TORCH_2_0 = check_version(TORCH_VERSION, "2.0.0")
+TORCH_2_1 = check_version(TORCH_VERSION, "2.1.0")
+TORCH_2_3 = check_version(TORCH_VERSION, "2.3.0")
+TORCH_2_4 = check_version(TORCH_VERSION, "2.4.0")
+TORCH_2_8 = check_version(TORCH_VERSION, "2.8.0")
+TORCH_2_9 = check_version(TORCH_VERSION, "2.9.0")
+TORCH_2_10 = check_version(TORCH_VERSION, "2.10.0")
+TORCHVISION_0_10 = check_version(TORCHVISION_VERSION, "0.10.0")
+TORCHVISION_0_11 = check_version(TORCHVISION_VERSION, "0.11.0")
+TORCHVISION_0_13 = check_version(TORCHVISION_VERSION, "0.13.0")
+TORCHVISION_0_18 = check_version(TORCHVISION_VERSION, "0.18.0")
+if WINDOWS and check_version(TORCH_VERSION, "==2.4.0"):  # reject version 2.4.0 on Windows
+    LOGGER.warning(
+        "Known issue with torch==2.4.0 on Windows with CPU, recommend upgrading to torch>=2.4.1 to resolve "
+        "https://github.com/ultralytics/ultralytics/issues/15049"
+    )
+
+
+@contextmanager
+def torch_distributed_zero_first(local_rank: int):
+    """Ensure all processes in distributed training wait for the local master (rank 0) to complete a task first."""
+    initialized = dist.is_available() and dist.is_initialized()
+    use_ids = initialized and dist.get_backend() == "nccl"
+
+    if initialized and local_rank not in {-1, 0}:
+        dist.barrier(device_ids=[local_rank]) if use_ids else dist.barrier()
+    yield
+    if initialized and local_rank == 0:
+        dist.barrier(device_ids=[local_rank]) if use_ids else dist.barrier()
+
+
+def smart_inference_mode():
+    """Apply torch.inference_mode() decorator if torch>=1.10.0, else torch.no_grad() decorator."""
+
+    def decorate(fn):
+        """Apply appropriate torch decorator for inference mode based on torch version."""
+        if TORCH_1_9 and torch.is_inference_mode_enabled():
+            return fn  # already in inference_mode, act as a pass-through
+        else:
+            return (torch.inference_mode if TORCH_1_10 else torch.no_grad)()(fn)
+
+    return decorate
+
+
+def autocast(enabled: bool, device: str = "cuda"):
+    """Get the appropriate autocast context manager based on PyTorch version and AMP setting.
+
+    This function returns a context manager for automatic mixed precision (AMP) training that is compatible with both
+    older and newer versions of PyTorch. It handles the differences in the autocast API between PyTorch versions.
+
+    Args:
+        enabled (bool): Whether to enable automatic mixed precision.
+        device (str, optional): The device to use for autocast.
+
+    Returns:
+        (torch.amp.autocast): The appropriate autocast context manager.
+
+    Examples:
+        >>> with autocast(enabled=True):
+        ...     # Your mixed precision operations here
+        ...     pass
+
+    Notes:
+        - For PyTorch versions 1.13 and newer, it uses `torch.amp.autocast`.
+        - For older versions, it uses `torch.cuda.amp.autocast`.
+    """
+    if TORCH_1_13:
+        return torch.amp.autocast(device, enabled=enabled)
+    else:
+        return torch.cuda.amp.autocast(enabled)
+
+
+@functools.lru_cache
+def get_cpu_info():
+    """Return a string with system CPU information, i.e. 'Apple M2'."""
+    from ultralytics.utils import PERSISTENT_CACHE  # avoid circular import error
+
+    if "cpu_info" not in PERSISTENT_CACHE:
+        try:
+            PERSISTENT_CACHE["cpu_info"] = CPUInfo.name()
+        except Exception:
+            pass
+    return PERSISTENT_CACHE.get("cpu_info", "unknown")
+
+
+@functools.lru_cache
+def get_gpu_info(index):
+    """Return a string with system GPU information, i.e. 'Tesla T4, 15102MiB'."""
+    properties = torch.cuda.get_device_properties(index)
+    return f"{properties.name}, {properties.total_memory / (1 << 20):.0f}MiB"
+
+
+def select_device(device="", newline=False, verbose=True):
+    """Select the appropriate PyTorch device based on the provided arguments.
+
+    The function takes a string specifying the device or a torch.device object and returns a torch.device object
+    representing the selected device. The function also validates the number of available devices and raises an
+    exception if the requested device(s) are not available.
+
+    Args:
+        device (str | torch.device, optional): Device string or torch.device object. Options include 'cpu', 'cuda', '0',
+            '0,1,2,3', 'mps', or '-1' for auto-select. Defaults to auto-selecting the first available GPU, or CPU if no
+            GPU is available.
+        newline (bool, optional): If True, adds a newline at the end of the log string.
+        verbose (bool, optional): If True, logs the device information.
+
+    Returns:
+        (torch.device): Selected device.
+
+    Examples:
+        >>> select_device("cuda:0")
+        device(type='cuda', index=0)
+
+        >>> select_device("cpu")
+        device(type='cpu')
+
+    Notes:
+        Sets the 'CUDA_VISIBLE_DEVICES' environment variable for specifying which GPUs to use.
+    """
+    if isinstance(device, torch.device) or str(device).startswith(("tpu", "intel", "vulkan")):
+        return device
+
+    s = f"Ultralytics {__version__} 🚀 Python-{PYTHON_VERSION} torch-{TORCH_VERSION} "
+    device = str(device).lower()
+    for remove in "cuda:", "none", "(", ")", "[", "]", "'", " ":
+        device = device.replace(remove, "")  # to string, 'cuda:0' -> '0' and '(0, 1)' -> '0,1'
+
+    # Auto-select GPUs
+    if "-1" in device:
+        from ultralytics.utils.autodevice import GPUInfo
+
+        # Replace each -1 with a selected GPU or remove it
+        parts = device.split(",")
+        selected = GPUInfo().select_idle_gpu(count=parts.count("-1"), min_memory_fraction=0.2)
+        for i in range(len(parts)):
+            if parts[i] == "-1":
+                parts[i] = str(selected.pop(0)) if selected else ""
+        device = ",".join(p for p in parts if p)
+
+    cpu = device == "cpu"
+    mps = device in {"mps", "mps:0"}  # Apple Metal Performance Shaders (MPS)
+    if cpu or mps:
+        os.environ["CUDA_VISIBLE_DEVICES"] = ""  # force torch.cuda.is_available() = False
+    elif device:  # non-cpu device requested
+        if device == "cuda":
+            device = "0"
+        if "," in device:
+            device = ",".join([x for x in device.split(",") if x])  # remove sequential commas, i.e. "0,,1" -> "0,1"
+        visible = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+        os.environ["CUDA_VISIBLE_DEVICES"] = device  # set environment variable - must be before assert is_available()
+        if not (torch.cuda.is_available() and torch.cuda.device_count() >= len(device.split(","))):
+            LOGGER.info(s)
+            install = (
+                "See https://pytorch.org/get-started/locally/ for up-to-date torch install instructions if no "
+                "CUDA devices are seen by torch.\n"
+                if torch.cuda.device_count() == 0
+                else ""
+            )
+            raise ValueError(
+                f"Invalid CUDA 'device={device}' requested."
+                f" Use 'device=cpu' or pass valid CUDA device(s) if available,"
+                f" i.e. 'device=0' or 'device=0,1,2,3' for Multi-GPU.\n"
+                f"\ntorch.cuda.is_available(): {torch.cuda.is_available()}"
+                f"\ntorch.cuda.device_count(): {torch.cuda.device_count()}"
+                f"\nos.environ['CUDA_VISIBLE_DEVICES']: {visible}\n"
+                f"{install}"
+            )
+
+    if not cpu and not mps and torch.cuda.is_available():  # prefer GPU if available
+        devices = device.split(",") if device else "0"  # i.e. "0,1" -> ["0", "1"]
+        space = " " * len(s)
+        for i, d in enumerate(devices):
+            s += f"{'' if i == 0 else space}CUDA:{d} ({get_gpu_info(i)})\n"  # bytes to MB
+        arg = "cuda:0"
+    elif mps and TORCH_2_0 and torch.backends.mps.is_available():
+        # Prefer MPS if available
+        s += f"MPS ({get_cpu_info()})\n"
+        arg = "mps"
+    else:  # revert to CPU
+        s += f"CPU ({get_cpu_info()})\n"
+        arg = "cpu"
+
+    if arg in {"cpu", "mps"}:
+        torch.set_num_threads(NUM_THREADS)  # reset OMP_NUM_THREADS for cpu training
+    if verbose:
+        LOGGER.info(s if newline else s.rstrip())
+    return torch.device(arg)
+
+
+def time_sync():
+    """Return PyTorch-accurate time."""
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    return time.time()
+
+
+def fuse_conv_and_bn(conv, bn):
+    """Fuse Conv2d and BatchNorm2d layers for inference optimization.
+
+    Args:
+        conv (nn.Conv2d): Convolutional layer to fuse.
+        bn (nn.BatchNorm2d): Batch normalization layer to fuse.
+
+    Returns:
+        (nn.Conv2d): The fused convolutional layer with gradients disabled.
+
+    Examples:
+        >>> conv = nn.Conv2d(3, 16, 3)
+        >>> bn = nn.BatchNorm2d(16)
+        >>> fused_conv = fuse_conv_and_bn(conv, bn)
+    """
+    # Compute fused weights
+    w_conv = conv.weight.view(conv.out_channels, -1)
+    w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
+    conv.weight.data = torch.mm(w_bn, w_conv).view(conv.weight.shape)
+
+    # Compute fused bias
+    b_conv = torch.zeros(conv.out_channels, device=conv.weight.device) if conv.bias is None else conv.bias
+    b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
+    fused_bias = torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn
+
+    if conv.bias is None:
+        conv.register_parameter("bias", nn.Parameter(fused_bias))
+    else:
+        conv.bias.data = fused_bias
+
+    return conv.requires_grad_(False)
+
+
+def fuse_deconv_and_bn(deconv, bn):
+    """Fuse ConvTranspose2d and BatchNorm2d layers for inference optimization.
+
+    Args:
+        deconv (nn.ConvTranspose2d): Transposed convolutional layer to fuse.
+        bn (nn.BatchNorm2d): Batch normalization layer to fuse.
+
+    Returns:
+        (nn.ConvTranspose2d): The fused transposed convolutional layer with gradients disabled.
+
+    Examples:
+        >>> deconv = nn.ConvTranspose2d(16, 3, 3)
+        >>> bn = nn.BatchNorm2d(3)
+        >>> fused_deconv = fuse_deconv_and_bn(deconv, bn)
+    """
+    # Compute fused weights
+    w_deconv = deconv.weight.view(deconv.out_channels, -1)
+    w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
+    deconv.weight.data = torch.mm(w_bn, w_deconv).view(deconv.weight.shape)
+
+    # Compute fused bias
+    b_conv = torch.zeros(deconv.out_channels, device=deconv.weight.device) if deconv.bias is None else deconv.bias
+    b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
+    fused_bias = torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn
+
+    if deconv.bias is None:
+        deconv.register_parameter("bias", nn.Parameter(fused_bias))
+    else:
+        deconv.bias.data = fused_bias
+
+    return deconv.requires_grad_(False)
+
+
+def model_info(model, detailed=False, verbose=True, imgsz=640):
+    """Print and return detailed model information layer by layer.
+
+    Args:
+        model (nn.Module): Model to analyze.
+        detailed (bool, optional): Whether to print detailed layer information.
+        verbose (bool, optional): Whether to print model information.
+        imgsz (int | list, optional): Input image size.
+
+    Returns:
+        (tuple): Tuple containing:
+            - n_l (int): Number of layers.
+            - n_p (int): Number of parameters.
+            - n_g (int): Number of gradients.
+            - flops (float): GFLOPs.
+    """
+    if not verbose:
+        return
+    n_p = get_num_params(model)  # number of parameters
+    n_g = get_num_gradients(model)  # number of gradients
+    layers = __import__("collections").OrderedDict((n, m) for n, m in model.named_modules() if len(m._modules) == 0)
+    n_l = len(layers)  # number of layers
+    if detailed:
+        h = f"{'layer':>5}{'name':>40}{'type':>20}{'gradient':>10}{'parameters':>12}{'shape':>20}{'mu':>10}{'sigma':>10}"
+        LOGGER.info(h)
+        for i, (mn, m) in enumerate(layers.items()):
+            mn = mn.replace("module_list.", "")
+            mt = m.__class__.__name__
+            if len(m._parameters):
+                for pn, p in m.named_parameters():
+                    LOGGER.info(
+                        f"{i:>5g}{f'{mn}.{pn}':>40}{mt:>20}{p.requires_grad!r:>10}{p.numel():>12g}{list(p.shape)!s:>20}{p.mean():>10.3g}{p.std():>10.3g}{str(p.dtype).replace('torch.', ''):>15}"
+                    )
+            else:  # layers with no learnable params
+                LOGGER.info(f"{i:>5g}{mn:>40}{mt:>20}{False!r:>10}{0:>12g}{[]!s:>20}{'-':>10}{'-':>10}{'-':>15}")
+
+    flops = get_flops(model, imgsz)  # imgsz may be int or list, i.e. imgsz=640 or imgsz=[640, 320]
+    fused = " (fused)" if getattr(model, "is_fused", lambda: False)() else ""
+    fs = f", {flops:.1f} GFLOPs" if flops else ""
+    yaml_file = getattr(model, "yaml_file", "") or getattr(model, "yaml", {}).get("yaml_file", "")
+    model_name = Path(yaml_file).stem.replace("yolo", "YOLO") or "Model"
+    LOGGER.info(f"{model_name} summary{fused}: {n_l:,} layers, {n_p:,} parameters, {n_g:,} gradients{fs}")
+    return n_l, n_p, n_g, flops
+
+
+def get_num_params(model):
+    """Return the total number of parameters in a YOLO model."""
+    return sum(x.numel() for x in model.parameters())
+
+
+def get_num_gradients(model):
+    """Return the total number of parameters with gradients in a YOLO model."""
+    return sum(x.numel() for x in model.parameters() if x.requires_grad)
+
+
+def model_info_for_loggers(trainer):
+    """Return model info dict with useful model information.
+
+    Args:
+        trainer (ultralytics.engine.trainer.BaseTrainer): The trainer object containing model and validation data.
+
+    Returns:
+        (dict): Dictionary containing model parameters, GFLOPs, and inference speeds.
+
+    Examples:
+        YOLOv8n info for loggers
+        >>> results = {
+        ...    "model/parameters": 3151904,
+        ...    "model/GFLOPs": 8.746,
+        ...    "model/speed_ONNX(ms)": 41.244,
+        ...    "model/speed_TensorRT(ms)": 3.211,
+        ...    "model/speed_PyTorch(ms)": 18.755,
+        ...}
+    """
+    if trainer.args.profile:  # profile ONNX and TensorRT times
+        from ultralytics.utils.benchmarks import ProfileModels
+
+        results = ProfileModels([trainer.last], device=trainer.device).run()[0]
+        results.pop("model/name")
+    else:  # only return PyTorch times from most recent validation
+        results = {
+            "model/parameters": get_num_params(trainer.model),
+            "model/GFLOPs": round(get_flops(trainer.model), 3),
+        }
+    results["model/speed_PyTorch(ms)"] = round(trainer.validator.speed["inference"], 3)
+    return results
+
+
+def get_flops(model, imgsz=640):
+    """Calculate FLOPs (floating point operations) for a model in GFLOPs.
+
+    Attempts two calculation methods: first with a stride-based tensor for efficiency, then falls back to full image
+    size if needed (e.g., for RTDETR models). Returns 0.0 if thop library is unavailable or calculation fails.
+
+    Args:
+        model (nn.Module): The model to calculate FLOPs for.
+        imgsz (int | list, optional): Input image size.
+
+    Returns:
+        (float): The model's GFLOPs (billions of floating point operations).
+    """
+    try:
+        import thop
+    except ImportError:
+        thop = None  # conda support without 'ultralytics-thop' installed
+
+    if not thop:
+        return 0.0  # if not installed return 0.0 GFLOPs
+
+    try:
+        model = unwrap_model(model)
+        p = next(model.parameters())
+        if not isinstance(imgsz, list):
+            imgsz = [imgsz, imgsz]  # expand if int/float
+        try:
+            # Method 1: Use stride-based input tensor
+            stride = max(int(model.stride.max()), 32) if hasattr(model, "stride") else 32  # max stride
+            im = torch.empty((1, p.shape[1], stride, stride), device=p.device)  # input image in BCHW format
+            flops = thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1e9 * 2  # stride GFLOPs
+            return flops * imgsz[0] / stride * imgsz[1] / stride  # imgsz GFLOPs
+        except Exception:
+            # Method 2: Use actual image size (required for RTDETR models)
+            im = torch.empty((1, p.shape[1], *imgsz), device=p.device)  # input image in BCHW format
+            return thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1e9 * 2  # imgsz GFLOPs
+    except Exception:
+        return 0.0
+
+
+def get_flops_with_torch_profiler(model, imgsz=640):
+    """Compute model FLOPs using torch profiler (alternative to thop package, but 2-10x slower).
+
+    Args:
+        model (nn.Module): The model to calculate FLOPs for.
+        imgsz (int | list, optional): Input image size.
+
+    Returns:
+        (float): The model's GFLOPs (billions of floating point operations).
+    """
+    if not TORCH_2_0:  # torch profiler implemented in torch>=2.0
+        return 0.0
+    model = unwrap_model(model)
+    p = next(model.parameters())
+    if not isinstance(imgsz, list):
+        imgsz = [imgsz, imgsz]  # expand if int/float
+    try:
+        # Use stride size for input tensor
+        stride = (max(int(model.stride.max()), 32) if hasattr(model, "stride") else 32) * 2  # max stride
+        im = torch.empty((1, p.shape[1], stride, stride), device=p.device)  # input image in BCHW format
+        with torch.profiler.profile(with_flops=True) as prof:
+            model(im)
+        flops = sum(x.flops for x in prof.key_averages()) / 1e9
+        flops = flops * imgsz[0] / stride * imgsz[1] / stride  # 640x640 GFLOPs
+    except Exception:
+        # Use actual image size for input tensor (i.e. required for RTDETR models)
+        im = torch.empty((1, p.shape[1], *imgsz), device=p.device)  # input image in BCHW format
+        with torch.profiler.profile(with_flops=True) as prof:
+            model(im)
+        flops = sum(x.flops for x in prof.key_averages()) / 1e9
+    return flops
+
+
+def initialize_weights(model):
+    """Initialize model weights, biases, and module settings to default values."""
+    for m in model.modules():
+        t = type(m)
+        if t is nn.Conv2d:
+            pass  # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+        elif t is nn.BatchNorm2d:
+            m.eps = 1e-3
+            m.momentum = 0.03
+        elif t in {nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU}:
+            m.inplace = True
+
+
+def scale_img(img, ratio=1.0, same_shape=False, gs=32):
+    """Scale and pad an image tensor, optionally maintaining aspect ratio and padding to gs multiple.
+
+    Args:
+        img (torch.Tensor): Input image tensor.
+        ratio (float, optional): Scaling ratio.
+        same_shape (bool, optional): Whether to maintain the same shape.
+        gs (int, optional): Grid size for padding.
+
+    Returns:
+        (torch.Tensor): Scaled and padded image tensor.
+    """
+    if ratio == 1.0:
+        return img
+    h, w = img.shape[2:]
+    s = (int(h * ratio), int(w * ratio))  # new size
+    img = F.interpolate(img, size=s, mode="bilinear", align_corners=False)  # resize
+    if not same_shape:  # pad/crop img
+        h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
+    return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)  # value = imagenet mean
+
+
+def copy_attr(a, b, include=(), exclude=()):
+    """Copy attributes from object 'b' to object 'a', with options to include/exclude certain attributes.
+
+    Args:
+        a (Any): Destination object to copy attributes to.
+        b (Any): Source object to copy attributes from.
+        include (tuple, optional): Attributes to include. If empty, all attributes are included.
+        exclude (tuple, optional): Attributes to exclude.
+    """
+    for k, v in b.__dict__.items():
+        if (len(include) and k not in include) or k.startswith("_") or k in exclude:
+            continue
+        else:
+            setattr(a, k, v)
+
+
+def intersect_dicts(da, db, exclude=()):
+    """Return a dictionary of intersecting keys with matching shapes, excluding 'exclude' keys, using da values.
+
+    Args:
+        da (dict): First dictionary.
+        db (dict): Second dictionary.
+        exclude (tuple, optional): Keys to exclude.
+
+    Returns:
+        (dict): Dictionary of intersecting keys with matching shapes.
+    """
+    return {k: v for k, v in da.items() if k in db and all(x not in k for x in exclude) and v.shape == db[k].shape}
+
+
+def is_parallel(model):
+    """Return True if model is of type DP or DDP.
+
+    Args:
+        model (nn.Module): Model to check.
+
+    Returns:
+        (bool): True if model is DataParallel or DistributedDataParallel.
+    """
+    return isinstance(model, (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel))
+
+
+def unwrap_model(m: nn.Module) -> nn.Module:
+    """Unwrap compiled and parallel models to get the base model.
+
+    Args:
+        m (nn.Module): A model that may be wrapped by torch.compile (._orig_mod) or parallel wrappers such as
+            DataParallel/DistributedDataParallel (.module).
+
+    Returns:
+        (nn.Module): The unwrapped base model without compile or parallel wrappers.
+    """
+    while True:
+        if hasattr(m, "_orig_mod") and isinstance(m._orig_mod, nn.Module):
+            m = m._orig_mod
+        elif hasattr(m, "module") and isinstance(m.module, nn.Module):
+            m = m.module
+        else:
+            return m
+
+
+def one_cycle(y1=0.0, y2=1.0, steps=100):
+    """Return a lambda function for sinusoidal ramp from y1 to y2 https://arxiv.org/pdf/1812.01187.pdf.
+
+    Args:
+        y1 (float, optional): Initial value.
+        y2 (float, optional): Final value.
+        steps (int, optional): Number of steps.
+
+    Returns:
+        (function): Lambda function for computing the sinusoidal ramp.
+    """
+    return lambda x: max((1 - math.cos(x * math.pi / steps)) / 2, 0) * (y2 - y1) + y1
+
+
+def init_seeds(seed=0, deterministic=False):
+    """Initialize random number generator (RNG) seeds https://pytorch.org/docs/stable/notes/randomness.html.
+
+    Args:
+        seed (int, optional): Random seed.
+        deterministic (bool, optional): Whether to set deterministic algorithms.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # for Multi-GPU, exception safe
+    # torch.backends.cudnn.benchmark = True  # AutoBatch problem https://github.com/ultralytics/yolov5/issues/9287
+    if deterministic:
+        if TORCH_2_0:
+            torch.use_deterministic_algorithms(True, warn_only=True)  # warn if deterministic is not possible
+            torch.backends.cudnn.deterministic = True
+            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+            os.environ["PYTHONHASHSEED"] = str(seed)
+        else:
+            LOGGER.warning("Upgrade to torch>=2.0.0 for deterministic training.")
+    else:
+        unset_deterministic()
+
+
+def unset_deterministic():
+    """Unset all the configurations applied for deterministic training."""
+    torch.use_deterministic_algorithms(False)
+    torch.backends.cudnn.deterministic = False
+    os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None)
+    os.environ.pop("PYTHONHASHSEED", None)
+
+
+class ModelEMA:
+    """Updated Exponential Moving Average (EMA) implementation.
+
+    Keeps a moving average of everything in the model state_dict (parameters and buffers). For EMA details see
+    References.
+
+    To disable EMA set the `enabled` attribute to `False`.
+
+    Attributes:
+        ema (nn.Module): Copy of the model in evaluation mode.
+        updates (int): Number of EMA updates.
+        decay (function): Decay function that determines the EMA weight.
+        enabled (bool): Whether EMA is enabled.
+
+    References:
+        - https://github.com/rwightman/pytorch-image-models
+        - https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
+    """
+
+    def __init__(self, model, decay=0.9999, tau=2000, updates=0):
+        """Initialize EMA for 'model' with given arguments.
+
+        Args:
+            model (nn.Module): Model to create EMA for.
+            decay (float, optional): Maximum EMA decay rate.
+            tau (int, optional): EMA decay time constant.
+            updates (int, optional): Initial number of updates.
+        """
+        self.ema = deepcopy(unwrap_model(model)).eval()  # FP32 EMA
+        self.updates = updates  # number of EMA updates
+        self.decay = lambda x: decay * (1 - math.exp(-x / tau))  # decay exponential ramp (to help early epochs)
+        for p in self.ema.parameters():
+            p.requires_grad_(False)
+        self.enabled = True
+
+    def update(self, model):
+        """Update EMA parameters.
+
+        Args:
+            model (nn.Module): Model to update EMA from.
+        """
+        if self.enabled:
+            self.updates += 1
+            d = self.decay(self.updates)
+
+            msd = unwrap_model(model).state_dict()  # model state_dict
+            for k, v in self.ema.state_dict().items():
+                if v.dtype.is_floating_point:  # true for FP16 and FP32
+                    v *= d
+                    v += (1 - d) * msd[k].detach()
+                    # assert v.dtype == msd[k].dtype == torch.float32, f'{k}: EMA {v.dtype},  model {msd[k].dtype}'
+
+    def update_attr(self, model, include=(), exclude=("process_group", "reducer")):
+        """Copy attributes from model to EMA, with options to include/exclude certain attributes.
+
+        Args:
+            model (nn.Module): Model to copy attributes from.
+            include (tuple, optional): Attributes to include.
+            exclude (tuple, optional): Attributes to exclude.
+        """
+        if self.enabled:
+            copy_attr(self.ema, model, include, exclude)
+
+
+def strip_optimizer(f: str | Path = "best.pt", s: str = "", updates: dict[str, Any] | None = None) -> dict[str, Any]:
+    """Strip optimizer from 'f' to finalize training, optionally save as 's'.
+
+    Args:
+        f (str | Path): File path to model to strip the optimizer from.
+        s (str, optional): File path to save the model with stripped optimizer to. If not provided, 'f' will be
+            overwritten.
+        updates (dict, optional): A dictionary of updates to overlay onto the checkpoint before saving.
+
+    Returns:
+        (dict): The combined checkpoint dictionary.
+
+    Examples:
+        >>> from pathlib import Path
+        >>> from ultralytics.utils.torch_utils import strip_optimizer
+        >>> for f in Path("path/to/model/checkpoints").rglob("*.pt"):
+        ...     strip_optimizer(f)
+    """
+    try:
+        x = torch_load(f, map_location=torch.device("cpu"))
+        assert isinstance(x, dict), "checkpoint is not a Python dictionary"
+        assert "model" in x, "'model' missing from checkpoint"
+    except Exception as e:
+        LOGGER.warning(f"Skipping {f}, not a valid Ultralytics model: {e}")
+        return {}
+
+    metadata = {
+        "date": datetime.now().isoformat(),
+        "version": __version__,
+        "license": "AGPL-3.0 License (https://ultralytics.com/license)",
+        "docs": "https://docs.ultralytics.com",
+    }
+
+    # Update model
+    if x.get("ema"):
+        x["model"] = x["ema"]  # replace model with EMA
+    if hasattr(x["model"], "args"):
+        x["model"].args = dict(x["model"].args)  # convert from IterableSimpleNamespace to dict
+    if hasattr(x["model"], "criterion"):
+        x["model"].criterion = None  # strip loss criterion
+    # x["model"].half()  # to FP16
+    for p in x["model"].parameters():
+        p.requires_grad = False
+
+    # Update other keys
+    args = {**DEFAULT_CFG_DICT, **x.get("train_args", {})}  # combine args
+    for k in "optimizer", "best_fitness", "ema", "updates", "scaler":  # keys
+        x[k] = None
+    x["epoch"] = -1
+    x["train_args"] = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS}  # strip non-default keys
+    # x['model'].args = x['train_args']
+
+    # Save
+    combined = {**metadata, **x, **(updates or {})}
+    torch.save(combined, s or f)  # combine dicts (prefer to the right)
+    mb = os.path.getsize(s or f) / 1e6  # file size
+    LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
+    return combined
+
+
+def convert_optimizer_state_dict_to_fp16(state_dict):
+    """Convert the state_dict of a given optimizer to FP16, focusing on the 'state' key for tensor conversions.
+
+    Args:
+        state_dict (dict): Optimizer state dictionary.
+
+    Returns:
+        (dict): Converted optimizer state dictionary with FP16 tensors.
+    """
+    for state in state_dict["state"].values():
+        for k, v in state.items():
+            if k != "step" and isinstance(v, torch.Tensor) and v.dtype is torch.float32:
+                state[k] = v.half()
+
+    return state_dict
+
+
+@contextmanager
+def cuda_memory_usage(device=None):
+    """Monitor and manage CUDA memory usage.
+
+    This function checks if CUDA is available and, if so, empties the CUDA cache to free up unused memory. It then
+    yields a dictionary containing memory usage information, which can be updated by the caller. Finally, it updates the
+    dictionary with the amount of memory reserved by CUDA on the specified device.
+
+    Args:
+        device (torch.device, optional): The CUDA device to query memory usage for.
+
+    Yields:
+        (dict): A dictionary with a key 'memory' initialized to 0, which will be updated with the reserved memory.
+    """
+    cuda_info = dict(memory=0)
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        try:
+            yield cuda_info
+        finally:
+            cuda_info["memory"] = torch.cuda.memory_reserved(device)
+    else:
+        yield cuda_info
+
+
+def profile_ops(input, ops, n=10, device=None, max_num_obj=0):
+    """Ultralytics speed, memory and FLOPs profiler.
+
+    Args:
+        input (torch.Tensor | list): Input tensor(s) to profile.
+        ops (nn.Module | list): Model or list of operations to profile.
+        n (int, optional): Number of iterations to average.
+        device (str | torch.device, optional): Device to profile on.
+        max_num_obj (int, optional): Maximum number of objects for simulation.
+
+    Returns:
+        (list): Profile results for each operation.
+
+    Examples:
+        >>> from ultralytics.utils.torch_utils import profile_ops
+        >>> input = torch.randn(16, 3, 640, 640)
+        >>> m1 = lambda x: x * torch.sigmoid(x)
+        >>> m2 = nn.SiLU()
+        >>> profile_ops(input, [m1, m2], n=100)  # profile over 100 iterations
+    """
+    try:
+        import thop
+    except ImportError:
+        thop = None  # conda support without 'ultralytics-thop' installed
+
+    results = []
+    if not isinstance(device, torch.device):
+        device = select_device(device)
+    LOGGER.info(
+        f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
+        f"{'input':>24s}{'output':>24s}"
+    )
+    gc.collect()  # attempt to free unused memory
+    torch.cuda.empty_cache()
+    for x in input if isinstance(input, list) else [input]:
+        x = x.to(device)
+        x.requires_grad = True
+        for m in ops if isinstance(ops, list) else [ops]:
+            m = m.to(device) if hasattr(m, "to") else m  # device
+            m = m.half() if hasattr(m, "half") and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m
+            tf, tb, t = 0, 0, [0, 0, 0]  # dt forward, backward
+            try:
+                flops = thop.profile(deepcopy(m), inputs=[x], verbose=False)[0] / 1e9 * 2 if thop else 0  # GFLOPs
+            except Exception:
+                flops = 0
+
+            try:
+                mem = 0
+                for _ in range(n):
+                    with cuda_memory_usage(device) as cuda_info:
+                        t[0] = time_sync()
+                        y = m(x)
+                        t[1] = time_sync()
+                        try:
+                            (sum(yi.sum() for yi in y) if isinstance(y, list) else y).sum().backward()
+                            t[2] = time_sync()
+                        except Exception:  # no backward method
+                            # print(e)  # for debug
+                            t[2] = float("nan")
+                    mem += cuda_info["memory"] / 1e9  # (GB)
+                    tf += (t[1] - t[0]) * 1000 / n  # ms per op forward
+                    tb += (t[2] - t[1]) * 1000 / n  # ms per op backward
+                    if max_num_obj:  # simulate training with predictions per image grid (for AutoBatch)
+                        with cuda_memory_usage(device) as cuda_info:
+                            torch.randn(
+                                x.shape[0],
+                                max_num_obj,
+                                int(sum((x.shape[-1] / s) * (x.shape[-2] / s) for s in m.stride.tolist())),
+                                device=device,
+                                dtype=torch.float32,
+                            )
+                        mem += cuda_info["memory"] / 1e9  # (GB)
+                s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else "list" for x in (x, y))  # shapes
+                p = sum(x.numel() for x in m.parameters()) if isinstance(m, nn.Module) else 0  # parameters
+                LOGGER.info(f"{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{s_in!s:>24s}{s_out!s:>24s}")
+                results.append([p, flops, mem, tf, tb, s_in, s_out])
+            except Exception as e:
+                LOGGER.info(e)
+                results.append(None)
+            finally:
+                gc.collect()  # attempt to free unused memory
+                torch.cuda.empty_cache()
+    return results
+
+
+class EarlyStopping:
+    """Early stopping class that stops training when a specified number of epochs have passed without improvement.
+
+    Attributes:
+        best_fitness (float): Best fitness value observed.
+        best_epoch (int): Epoch where best fitness was observed.
+        patience (int): Number of epochs to wait after fitness stops improving before stopping.
+        possible_stop (bool): Flag indicating if stopping may occur next epoch.
+    """
+
+    def __init__(self, patience=50):
+        """Initialize early stopping object.
+
+        Args:
+            patience (int, optional): Number of epochs to wait after fitness stops improving before stopping.
+        """
+        self.best_fitness = 0.0  # i.e. mAP
+        self.best_epoch = 0
+        self.patience = patience or float("inf")  # epochs to wait after fitness stops improving to stop
+        self.possible_stop = False  # possible stop may occur next epoch
+
+    def __call__(self, epoch, fitness):
+        """Check whether to stop training.
+
+        Args:
+            epoch (int): Current epoch of training.
+            fitness (float): Fitness value of current epoch.
+
+        Returns:
+            (bool): True if training should stop, False otherwise.
+        """
+        if fitness is None:  # check if fitness=None (happens when val=False)
+            return False
+
+        if fitness > self.best_fitness or self.best_fitness == 0:  # allow for early zero-fitness stage of training
+            self.best_epoch = epoch
+            self.best_fitness = fitness
+        delta = epoch - self.best_epoch  # epochs without improvement
+        self.possible_stop = delta >= (self.patience - 1)  # possible stop may occur next epoch
+        stop = delta >= self.patience  # stop training if patience exceeded
+        if stop:
+            prefix = colorstr("EarlyStopping: ")
+            LOGGER.info(
+                f"{prefix}Training stopped early as no improvement observed in last {self.patience} epochs. "
+                f"Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n"
+                f"To update EarlyStopping(patience={self.patience}) pass a new patience value, "
+                f"i.e. `patience=300` or use `patience=0` to disable EarlyStopping."
+            )
+        return stop
+
+
+def attempt_compile(
+    model: torch.nn.Module,
+    device: torch.device,
+    imgsz: int = 640,
+    use_autocast: bool = False,
+    warmup: bool = False,
+    mode: bool | str = "default",
+) -> torch.nn.Module:
+    """Compile a model with torch.compile and optionally warm up the graph to reduce first-iteration latency.
+
+    This utility attempts to compile the provided model using the inductor backend. If compilation is unavailable or
+    fails, the original model is returned unchanged. An optional warmup performs a single forward pass on a dummy input
+    to prime the compiled graph and measure compile/warmup time.
+
+    Args:
+        model (torch.nn.Module): Model to compile.
+        device (torch.device): Inference device used for warmup and autocast decisions.
+        imgsz (int, optional): Square input size to create a dummy tensor with shape (1, 3, imgsz, imgsz) for warmup.
+        use_autocast (bool, optional): Whether to run warmup under autocast on CUDA or MPS devices.
+        warmup (bool, optional): Whether to execute a single dummy forward pass to warm up the compiled model.
+        mode (bool | str, optional): torch.compile mode. True → "default", False → no compile, or a string like
+            "default", "reduce-overhead", "max-autotune-no-cudagraphs".
+
+    Returns:
+        (torch.nn.Module): Compiled model if compilation succeeds, otherwise the original unmodified model.
+
+    Examples:
+        >>> device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        >>> # Try to compile and warm up a model with a 640x640 input
+        >>> model = attempt_compile(model, device=device, imgsz=640, use_autocast=True, warmup=True)
+
+    Notes:
+        - If the current PyTorch build does not provide torch.compile, the function returns the input model immediately.
+        - Warmup runs under torch.inference_mode and may use torch.autocast for CUDA/MPS to align compute precision.
+        - CUDA devices are synchronized after warmup to account for asynchronous kernel execution.
+    """
+    if not hasattr(torch, "compile") or not mode:
+        return model
+
+    if mode is True:
+        mode = "default"
+    prefix = colorstr("compile:")
+    LOGGER.info(f"{prefix} starting torch.compile with '{mode}' mode...")
+    if mode == "max-autotune":
+        LOGGER.warning(f"{prefix} mode='{mode}' not recommended, using mode='max-autotune-no-cudagraphs' instead")
+        mode = "max-autotune-no-cudagraphs"
+    t0 = time.perf_counter()
+    try:
+        model = torch.compile(model, mode=mode, backend="inductor")
+    except Exception as e:
+        LOGGER.warning(f"{prefix} torch.compile failed, continuing uncompiled: {e}")
+        return model
+    t_compile = time.perf_counter() - t0
+
+    t_warm = 0.0
+    if warmup:
+        # Use a single dummy tensor to build the graph shape state and reduce first-iteration latency
+        dummy = torch.zeros(1, 3, imgsz, imgsz, device=device)
+        if use_autocast and device.type == "cuda":
+            dummy = dummy.half()
+        t1 = time.perf_counter()
+        with torch.inference_mode():
+            if use_autocast and device.type in {"cuda", "mps"}:
+                with torch.autocast(device.type):
+                    _ = model(dummy)
+            else:
+                _ = model(dummy)
+        if device.type == "cuda":
+            torch.cuda.synchronize(device)
+        t_warm = time.perf_counter() - t1
+
+    total = t_compile + t_warm
+    if warmup:
+        LOGGER.info(f"{prefix} complete in {total:.1f}s (compile {t_compile:.1f}s + warmup {t_warm:.1f}s)")
+    else:
+        LOGGER.info(f"{prefix} compile complete in {t_compile:.1f}s (no warmup)")
+    return model

ultralytics/utils/tqdm.py

@@ -0,0 +1,443 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import os
+import sys
+import time
+from functools import lru_cache
+from typing import IO, Any
+
+
+@lru_cache(maxsize=1)
+def is_noninteractive_console() -> bool:
+    """Check for known non-interactive console environments."""
+    return "GITHUB_ACTIONS" in os.environ or "RUNPOD_POD_ID" in os.environ
+
+
+class TQDM:
+    """Lightweight zero-dependency progress bar for Ultralytics.
+
+    Provides clean, rich-style progress bars suitable for various environments including Weights & Biases, console
+    outputs, and other logging systems. Features zero external dependencies, clean single-line output, rich-style
+    progress bars with Unicode block characters, context manager support, iterator protocol support, and dynamic
+    description updates.
+
+    Attributes:
+        iterable (Any): Iterable to wrap with progress bar.
+        desc (str): Prefix description for the progress bar.
+        total (int | None): Expected number of iterations.
+        disable (bool): Whether to disable the progress bar.
+        unit (str): String for units of iteration.
+        unit_scale (bool): Auto-scale units flag.
+        unit_divisor (int): Divisor for unit scaling.
+        leave (bool): Whether to leave the progress bar after completion.
+        mininterval (float): Minimum time interval between updates.
+        initial (int): Initial counter value.
+        n (int): Current iteration count.
+        closed (bool): Whether the progress bar is closed.
+        bar_format (str | None): Custom bar format string.
+        file (IO[str]): Output file stream.
+
+    Methods:
+        update: Update progress by n steps.
+        set_description: Set or update the description.
+        set_postfix: Set postfix for the progress bar.
+        close: Close the progress bar and clean up.
+        refresh: Refresh the progress bar display.
+        clear: Clear the progress bar from display.
+        write: Write a message without breaking the progress bar.
+
+    Examples:
+        Basic usage with iterator:
+        >>> for i in TQDM(range(100)):
+        ...     time.sleep(0.01)
+
+        With custom description:
+        >>> pbar = TQDM(range(100), desc="Processing")
+        >>> for i in pbar:
+        ...     pbar.set_description(f"Processing item {i}")
+
+        Context manager usage:
+        >>> with TQDM(total=100, unit="B", unit_scale=True) as pbar:
+        ...     for i in range(100):
+        ...         pbar.update(1)
+
+        Manual updates:
+        >>> pbar = TQDM(total=100, desc="Training")
+        >>> for epoch in range(100):
+        ...     # Do work
+        ...     pbar.update(1)
+        >>> pbar.close()
+    """
+
+    # Constants
+    MIN_RATE_CALC_INTERVAL = 0.01  # Minimum time interval for rate calculation
+    RATE_SMOOTHING_FACTOR = 0.3  # Factor for exponential smoothing of rates
+    MAX_SMOOTHED_RATE = 1000000  # Maximum rate to apply smoothing to
+    NONINTERACTIVE_MIN_INTERVAL = 60.0  # Minimum interval for non-interactive environments
+
+    def __init__(
+        self,
+        iterable: Any = None,
+        desc: str | None = None,
+        total: int | None = None,
+        leave: bool = True,
+        file: IO[str] | None = None,
+        mininterval: float = 0.1,
+        disable: bool | None = None,
+        unit: str = "it",
+        unit_scale: bool = True,
+        unit_divisor: int = 1000,
+        bar_format: str | None = None,  # kept for API compatibility; not used for formatting
+        initial: int = 0,
+        **kwargs,
+    ) -> None:
+        """Initialize the TQDM progress bar with specified configuration options.
+
+        Args:
+            iterable (Any, optional): Iterable to wrap with progress bar.
+            desc (str, optional): Prefix description for the progress bar.
+            total (int, optional): Expected number of iterations.
+            leave (bool, optional): Whether to leave the progress bar after completion.
+            file (IO[str], optional): Output file stream for progress display.
+            mininterval (float, optional): Minimum time interval between updates (default 0.1s, 60s in GitHub Actions).
+            disable (bool, optional): Whether to disable the progress bar. Auto-detected if None.
+            unit (str, optional): String for units of iteration (default "it" for items).
+            unit_scale (bool, optional): Auto-scale units for bytes/data units.
+            unit_divisor (int, optional): Divisor for unit scaling (default 1000).
+            bar_format (str, optional): Custom bar format string.
+            initial (int, optional): Initial counter value.
+            **kwargs (Any): Additional keyword arguments for compatibility (ignored).
+        """
+        # Disable if not verbose
+        if disable is None:
+            try:
+                from ultralytics.utils import LOGGER, VERBOSE
+
+                disable = not VERBOSE or LOGGER.getEffectiveLevel() > 20
+            except ImportError:
+                disable = False
+
+        self.iterable = iterable
+        self.desc = desc or ""
+        self.total = total or (len(iterable) if hasattr(iterable, "__len__") else None) or None  # prevent total=0
+        self.disable = disable
+        self.unit = unit
+        self.unit_scale = unit_scale
+        self.unit_divisor = unit_divisor
+        self.leave = leave
+        self.noninteractive = is_noninteractive_console()
+        self.mininterval = max(mininterval, self.NONINTERACTIVE_MIN_INTERVAL) if self.noninteractive else mininterval
+        self.initial = initial
+
+        # Kept for API compatibility (unused for f-string formatting)
+        self.bar_format = bar_format
+
+        self.file = file or sys.stdout
+
+        # Internal state
+        self.n = self.initial
+        self.last_print_n = self.initial
+        self.last_print_t = time.time()
+        self.start_t = time.time()
+        self.last_rate = 0.0
+        self.closed = False
+        self.is_bytes = unit_scale and unit in {"B", "bytes"}
+        self.scales = (
+            [(1073741824, "GB/s"), (1048576, "MB/s"), (1024, "KB/s")]
+            if self.is_bytes
+            else [(1e9, f"G{self.unit}/s"), (1e6, f"M{self.unit}/s"), (1e3, f"K{self.unit}/s")]
+        )
+
+        if not self.disable and self.total and not self.noninteractive:
+            self._display()
+
+    def _format_rate(self, rate: float) -> str:
+        """Format rate with units, switching between it/s and s/it for readability."""
+        if rate <= 0:
+            return ""
+
+        inv_rate = 1 / rate if rate else None
+
+        # Use s/it format when inv_rate > 1 (i.e., rate < 1 it/s) for better readability
+        if inv_rate and inv_rate > 1:
+            return f"{inv_rate:.1f}s/B" if self.is_bytes else f"{inv_rate:.1f}s/{self.unit}"
+
+        # Use it/s format for fast iterations
+        fallback = f"{rate:.1f}B/s" if self.is_bytes else f"{rate:.1f}{self.unit}/s"
+        return next((f"{rate / t:.1f}{u}" for t, u in self.scales if rate >= t), fallback)
+
+    def _format_num(self, num: int | float) -> str:
+        """Format number with optional unit scaling."""
+        if not self.unit_scale or not self.is_bytes:
+            return str(num)
+
+        for unit in ("", "K", "M", "G", "T"):
+            if abs(num) < self.unit_divisor:
+                return f"{num:3.1f}{unit}B" if unit else f"{num:.0f}B"
+            num /= self.unit_divisor
+        return f"{num:.1f}PB"
+
+    @staticmethod
+    def _format_time(seconds: float) -> str:
+        """Format time duration."""
+        if seconds < 60:
+            return f"{seconds:.1f}s"
+        elif seconds < 3600:
+            return f"{int(seconds // 60)}:{seconds % 60:02.0f}"
+        else:
+            h, m = int(seconds // 3600), int((seconds % 3600) // 60)
+            return f"{h}:{m:02d}:{seconds % 60:02.0f}"
+
+    def _generate_bar(self, width: int = 12) -> str:
+        """Generate progress bar."""
+        if self.total is None:
+            return "━" * width if self.closed else "─" * width
+
+        frac = min(1.0, self.n / self.total)
+        filled = int(frac * width)
+        bar = "━" * filled + "─" * (width - filled)
+        if filled < width and frac * width - filled > 0.5:
+            bar = f"{bar[:filled]}╸{bar[filled + 1 :]}"
+        return bar
+
+    def _should_update(self, dt: float, dn: int) -> bool:
+        """Check if display should update."""
+        if self.noninteractive:
+            return False
+        return (self.total is not None and self.n >= self.total) or (dt >= self.mininterval)
+
+    def _display(self, final: bool = False) -> None:
+        """Display progress bar."""
+        if self.disable or (self.closed and not final):
+            return
+
+        current_time = time.time()
+        dt = current_time - self.last_print_t
+        dn = self.n - self.last_print_n
+
+        if not final and not self._should_update(dt, dn):
+            return
+
+        # Calculate rate (avoid crazy numbers)
+        if dt > self.MIN_RATE_CALC_INTERVAL:
+            rate = dn / dt if dt else 0.0
+            # Smooth rate for reasonable values, use raw rate for very high values
+            if rate < self.MAX_SMOOTHED_RATE:
+                self.last_rate = self.RATE_SMOOTHING_FACTOR * rate + (1 - self.RATE_SMOOTHING_FACTOR) * self.last_rate
+                rate = self.last_rate
+        else:
+            rate = self.last_rate
+
+        # At completion, use overall rate
+        if self.total and self.n >= self.total:
+            overall_elapsed = current_time - self.start_t
+            if overall_elapsed > 0:
+                rate = self.n / overall_elapsed
+
+        # Update counters
+        self.last_print_n = self.n
+        self.last_print_t = current_time
+        elapsed = current_time - self.start_t
+
+        # Remaining time
+        remaining_str = ""
+        if self.total and 0 < self.n < self.total and elapsed > 0:
+            est_rate = rate or (self.n / elapsed)
+            remaining_str = f"<{self._format_time((self.total - self.n) / est_rate)}"
+
+        # Numbers and percent
+        if self.total:
+            percent = (self.n / self.total) * 100
+            n_str = self._format_num(self.n)
+            t_str = self._format_num(self.total)
+            if self.is_bytes and n_str[-2] == t_str[-2]:  # Collapse suffix only when identical (e.g. "5.4/5.4MB")
+                n_str = n_str.rstrip("KMGTPB")
+        else:
+            percent = 0.0
+            n_str, t_str = self._format_num(self.n), "?"
+
+        elapsed_str = self._format_time(elapsed)
+        rate_str = self._format_rate(rate) or (self._format_rate(self.n / elapsed) if elapsed > 0 else "")
+
+        bar = self._generate_bar()
+
+        # Compose progress line via f-strings (two shapes: with/without total)
+        if self.total:
+            if self.is_bytes and self.n >= self.total:
+                # Completed bytes: show only final size
+                progress_str = f"{self.desc}: {percent:.0f}% {bar} {t_str} {rate_str} {elapsed_str}"
+            else:
+                progress_str = (
+                    f"{self.desc}: {percent:.0f}% {bar} {n_str}/{t_str} {rate_str} {elapsed_str}{remaining_str}"
+                )
+        else:
+            progress_str = f"{self.desc}: {bar} {n_str} {rate_str} {elapsed_str}"
+
+        # Write to output
+        try:
+            if self.noninteractive:
+                # In non-interactive environments, avoid carriage return which creates empty lines
+                self.file.write(progress_str)
+            else:
+                # In interactive terminals, use carriage return and clear line for updating display
+                self.file.write(f"\r\033[K{progress_str}")
+            self.file.flush()
+        except Exception:
+            pass
+
+    def update(self, n: int = 1) -> None:
+        """Update progress by n steps."""
+        if not self.disable and not self.closed:
+            self.n += n
+            self._display()
+
+    def set_description(self, desc: str | None) -> None:
+        """Set description."""
+        self.desc = desc or ""
+        if not self.disable:
+            self._display()
+
+    def set_postfix(self, **kwargs: Any) -> None:
+        """Set postfix (appends to description)."""
+        if kwargs:
+            postfix = ", ".join(f"{k}={v}" for k, v in kwargs.items())
+            base_desc = self.desc.split(" | ")[0] if " | " in self.desc else self.desc
+            self.set_description(f"{base_desc} | {postfix}")
+
+    def close(self) -> None:
+        """Close progress bar."""
+        if self.closed:
+            return
+
+        self.closed = True
+
+        if not self.disable:
+            # Final display
+            if self.total and self.n >= self.total:
+                self.n = self.total
+                if self.n != self.last_print_n:  # Skip if 100% already shown
+                    self._display(final=True)
+            else:
+                self._display(final=True)
+
+            # Cleanup
+            if self.leave:
+                self.file.write("\n")
+            else:
+                self.file.write("\r\033[K")
+
+            try:
+                self.file.flush()
+            except Exception:
+                pass
+
+    def __enter__(self) -> TQDM:
+        """Enter context manager."""
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        """Exit context manager and close progress bar."""
+        self.close()
+
+    def __iter__(self) -> Any:
+        """Iterate over the wrapped iterable with progress updates."""
+        if self.iterable is None:
+            raise TypeError("'NoneType' object is not iterable")
+
+        try:
+            for item in self.iterable:
+                yield item
+                self.update(1)
+        finally:
+            self.close()
+
+    def __del__(self) -> None:
+        """Destructor to ensure cleanup."""
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    def refresh(self) -> None:
+        """Refresh display."""
+        if not self.disable:
+            self._display()
+
+    def clear(self) -> None:
+        """Clear progress bar."""
+        if not self.disable:
+            try:
+                self.file.write("\r\033[K")
+                self.file.flush()
+            except Exception:
+                pass
+
+    @staticmethod
+    def write(s: str, file: IO[str] | None = None, end: str = "\n") -> None:
+        """Static method to write without breaking progress bar."""
+        file = file or sys.stdout
+        try:
+            file.write(s + end)
+            file.flush()
+        except Exception:
+            pass
+
+
+if __name__ == "__main__":
+    import time
+
+    print("1. Basic progress bar with known total:")
+    for i in TQDM(range(3), desc="Known total"):
+        time.sleep(0.05)
+
+    print("\n2. Manual updates with known total:")
+    pbar = TQDM(total=300, desc="Manual updates", unit="files")
+    for i in range(300):
+        time.sleep(0.03)
+        pbar.update(1)
+        if i % 10 == 9:
+            pbar.set_description(f"Processing batch {i // 10 + 1}")
+    pbar.close()
+
+    print("\n3. Progress bar with unknown total:")
+    pbar = TQDM(desc="Unknown total", unit="items")
+    for i in range(25):
+        time.sleep(0.08)
+        pbar.update(1)
+        if i % 5 == 4:
+            pbar.set_postfix(processed=i + 1, status="OK")
+    pbar.close()
+
+    print("\n4. Context manager with unknown total:")
+    with TQDM(desc="Processing stream", unit="B", unit_scale=True, unit_divisor=1024) as pbar:
+        for i in range(30):
+            time.sleep(0.1)
+            pbar.update(1024 * 1024 * i)  # Simulate processing MB of data
+
+    print("\n5. Iterator with unknown length:")
+
+    def data_stream():
+        """Simulate a data stream of unknown length."""
+        import random
+
+        for i in range(random.randint(10, 20)):
+            yield f"data_chunk_{i}"
+
+    for chunk in TQDM(data_stream(), desc="Stream processing", unit="chunks"):
+        time.sleep(0.1)
+
+    print("\n6. File processing simulation (unknown size):")
+
+    def process_files():
+        """Simulate processing files of unknown count."""
+        return [f"file_{i}.txt" for i in range(18)]
+
+    pbar = TQDM(desc="Scanning files", unit="files")
+    files = process_files()
+    for i, filename in enumerate(files):
+        time.sleep(0.06)
+        pbar.update(1)
+        pbar.set_description(f"Processing {filename}")
+    pbar.close()

ultralytics/utils/triton.py

@@ -0,0 +1,112 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+import ast
+from urllib.parse import urlsplit
+
+import numpy as np
+
+
+class TritonRemoteModel:
+    """Client for interacting with a remote Triton Inference Server model.
+
+    This class provides a convenient interface for sending inference requests to a Triton Inference Server and
+    processing the responses. Supports both HTTP and gRPC communication protocols.
+
+    Attributes:
+        endpoint (str): The name of the model on the Triton server.
+        url (str): The URL of the Triton server.
+        triton_client: The Triton client (either HTTP or gRPC).
+        InferInput: The input class for the Triton client.
+        InferRequestedOutput: The output request class for the Triton client.
+        input_formats (list[str]): The data types of the model inputs.
+        np_input_formats (list[type]): The numpy data types of the model inputs.
+        input_names (list[str]): The names of the model inputs.
+        output_names (list[str]): The names of the model outputs.
+        metadata: The metadata associated with the model.
+
+    Methods:
+        __call__: Call the model with the given inputs and return the outputs.
+
+    Examples:
+        Initialize a Triton client with HTTP
+        >>> model = TritonRemoteModel(url="localhost:8000", endpoint="yolov8", scheme="http")
+
+        Make inference with numpy arrays
+        >>> outputs = model(np.random.rand(1, 3, 640, 640).astype(np.float32))
+    """
+
+    def __init__(self, url: str, endpoint: str = "", scheme: str = ""):
+        """Initialize the TritonRemoteModel for interacting with a remote Triton Inference Server.
+
+        Arguments may be provided individually or parsed from a collective 'url' argument of the form
+        <scheme>://<netloc>/<endpoint>/<task_name>
+
+        Args:
+            url (str): The URL of the Triton server.
+            endpoint (str, optional): The name of the model on the Triton server.
+            scheme (str, optional): The communication scheme ('http' or 'grpc').
+        """
+        if not endpoint and not scheme:  # Parse all args from URL string
+            splits = urlsplit(url)
+            endpoint = splits.path.strip("/").split("/", 1)[0]
+            scheme = splits.scheme
+            url = splits.netloc
+
+        self.endpoint = endpoint
+        self.url = url
+
+        # Choose the Triton client based on the communication scheme
+        if scheme == "http":
+            import tritonclient.http as client
+
+            self.triton_client = client.InferenceServerClient(url=self.url, verbose=False, ssl=False)
+            config = self.triton_client.get_model_config(endpoint)
+        else:
+            import tritonclient.grpc as client
+
+            self.triton_client = client.InferenceServerClient(url=self.url, verbose=False, ssl=False)
+            config = self.triton_client.get_model_config(endpoint, as_json=True)["config"]
+
+        # Sort output names alphabetically, i.e. 'output0', 'output1', etc.
+        config["output"] = sorted(config["output"], key=lambda x: x.get("name"))
+
+        # Define model attributes
+        type_map = {"TYPE_FP32": np.float32, "TYPE_FP16": np.float16, "TYPE_UINT8": np.uint8}
+        self.InferRequestedOutput = client.InferRequestedOutput
+        self.InferInput = client.InferInput
+        self.input_formats = [x["data_type"] for x in config["input"]]
+        self.np_input_formats = [type_map[x] for x in self.input_formats]
+        self.input_names = [x["name"] for x in config["input"]]
+        self.output_names = [x["name"] for x in config["output"]]
+        self.metadata = ast.literal_eval(config.get("parameters", {}).get("metadata", {}).get("string_value", "None"))
+
+    def __call__(self, *inputs: np.ndarray) -> list[np.ndarray]:
+        """Call the model with the given inputs and return inference results.
+
+        Args:
+            *inputs (np.ndarray): Input data to the model. Each array should match the expected shape and type for the
+                corresponding model input.
+
+        Returns:
+            (list[np.ndarray]): Model outputs cast to the dtype of the first input. Each element in the list corresponds
+                to one of the model's output tensors.
+
+        Examples:
+            >>> model = TritonRemoteModel(url="localhost:8000", endpoint="yolov8", scheme="http")
+            >>> outputs = model(np.random.rand(1, 3, 640, 640).astype(np.float32))
+        """
+        infer_inputs = []
+        input_format = inputs[0].dtype
+        for i, x in enumerate(inputs):
+            if x.dtype != self.np_input_formats[i]:
+                x = x.astype(self.np_input_formats[i])
+            infer_input = self.InferInput(self.input_names[i], [*x.shape], self.input_formats[i].replace("TYPE_", ""))
+            infer_input.set_data_from_numpy(x)
+            infer_inputs.append(infer_input)
+
+        infer_outputs = [self.InferRequestedOutput(output_name) for output_name in self.output_names]
+        outputs = self.triton_client.infer(model_name=self.endpoint, inputs=infer_inputs, outputs=infer_outputs)
+
+        return [outputs.as_numpy(output_name).astype(input_format) for output_name in self.output_names]

ultralytics/utils/tuner.py

@@ -0,0 +1,168 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from ultralytics.cfg import TASK2DATA, TASK2METRIC, get_cfg, get_save_dir
+from ultralytics.utils import DEFAULT_CFG, DEFAULT_CFG_DICT, LOGGER, NUM_THREADS, checks, colorstr
+
+
+def run_ray_tune(
+    model,
+    space: dict | None = None,
+    grace_period: int = 10,
+    gpu_per_trial: int | None = None,
+    max_samples: int = 10,
+    **train_args,
+):
+    """Run hyperparameter tuning using Ray Tune.
+
+    Args:
+        model (YOLO): Model to run the tuner on.
+        space (dict, optional): The hyperparameter search space. If not provided, uses default space.
+        grace_period (int, optional): The grace period in epochs of the ASHA scheduler.
+        gpu_per_trial (int, optional): The number of GPUs to allocate per trial.
+        max_samples (int, optional): The maximum number of trials to run.
+        **train_args (Any): Additional arguments to pass to the `train()` method.
+
+    Returns:
+        (ray.tune.ResultGrid): A ResultGrid containing the results of the hyperparameter search.
+
+    Examples:
+        >>> from ultralytics import YOLO
+        >>> model = YOLO("yolo26n.pt")  # Load a YOLO26n model
+
+        Start tuning hyperparameters for YOLO26n training on the COCO8 dataset
+        >>> result_grid = model.tune(data="coco8.yaml", use_ray=True)
+    """
+    LOGGER.info("💡 Learn about RayTune at https://docs.ultralytics.com/integrations/ray-tune")
+    try:
+        checks.check_requirements("ray[tune]")
+
+        import ray
+        from ray import tune
+        from ray.air import RunConfig
+        from ray.air.integrations.wandb import WandbLoggerCallback
+        from ray.tune.schedulers import ASHAScheduler
+    except ImportError:
+        raise ModuleNotFoundError('Ray Tune required but not found. To install run: pip install "ray[tune]"')
+
+    try:
+        import wandb
+
+        assert hasattr(wandb, "__version__")
+    except (ImportError, AssertionError):
+        wandb = False
+
+    checks.check_version(ray.__version__, ">=2.0.0", "ray")
+    default_space = {
+        # 'optimizer': tune.choice(['SGD', 'Adam', 'AdamW', 'NAdam', 'RAdam', 'RMSProp']),
+        "lr0": tune.uniform(1e-5, 1e-1),
+        "lrf": tune.uniform(0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
+        "momentum": tune.uniform(0.6, 0.98),  # SGD momentum/Adam beta1
+        "weight_decay": tune.uniform(0.0, 0.001),  # optimizer weight decay
+        "warmup_epochs": tune.uniform(0.0, 5.0),  # warmup epochs (fractions ok)
+        "warmup_momentum": tune.uniform(0.0, 0.95),  # warmup initial momentum
+        "box": tune.uniform(0.02, 0.2),  # box loss gain
+        "cls": tune.uniform(0.2, 4.0),  # cls loss gain (scale with pixels)
+        "hsv_h": tune.uniform(0.0, 0.1),  # image HSV-Hue augmentation (fraction)
+        "hsv_s": tune.uniform(0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
+        "hsv_v": tune.uniform(0.0, 0.9),  # image HSV-Value augmentation (fraction)
+        "degrees": tune.uniform(0.0, 45.0),  # image rotation (+/- deg)
+        "translate": tune.uniform(0.0, 0.9),  # image translation (+/- fraction)
+        "scale": tune.uniform(0.0, 0.9),  # image scale (+/- gain)
+        "shear": tune.uniform(0.0, 10.0),  # image shear (+/- deg)
+        "perspective": tune.uniform(0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
+        "flipud": tune.uniform(0.0, 1.0),  # image flip up-down (probability)
+        "fliplr": tune.uniform(0.0, 1.0),  # image flip left-right (probability)
+        "bgr": tune.uniform(0.0, 1.0),  # swap RGB↔BGR channels (probability)
+        "mosaic": tune.uniform(0.0, 1.0),  # image mosaic (probability)
+        "mixup": tune.uniform(0.0, 1.0),  # image mixup (probability)
+        "cutmix": tune.uniform(0.0, 1.0),  # image cutmix (probability)
+        "copy_paste": tune.uniform(0.0, 1.0),  # segment copy-paste (probability)
+    }
+
+    # Put the model in ray store
+    task = model.task
+    model_in_store = ray.put(model)
+    base_name = train_args.get("name", "tune")
+
+    def _tune(config):
+        """Train the YOLO model with the specified hyperparameters and return results."""
+        model_to_train = ray.get(model_in_store)  # get the model from ray store for tuning
+        model_to_train.reset_callbacks()
+        config.update(train_args)
+
+        # Set trial-specific name for W&B logging
+        try:
+            trial_id = tune.get_trial_id()  # Get current trial ID (e.g., "2c2fc_00000")
+            trial_suffix = trial_id.split("_")[-1] if "_" in trial_id else trial_id
+            config["name"] = f"{base_name}_{trial_suffix}"
+        except Exception:
+            # Not in Ray Tune context or error getting trial ID, use base name
+            config["name"] = base_name
+
+        results = model_to_train.train(**config)
+        return results.results_dict
+
+    # Get search space
+    if not space and not train_args.get("resume"):
+        space = default_space
+        LOGGER.warning("Search space not provided, using default search space.")
+
+    # Get dataset
+    data = train_args.get("data", TASK2DATA[task])
+    space["data"] = data
+    if "data" not in train_args:
+        LOGGER.warning(f'Data not provided, using default "data={data}".')
+
+    # Define the trainable function with allocated resources
+    trainable_with_resources = tune.with_resources(_tune, {"cpu": NUM_THREADS, "gpu": gpu_per_trial or 0})
+
+    # Define the ASHA scheduler for hyperparameter search
+    asha_scheduler = ASHAScheduler(
+        time_attr="epoch",
+        metric=TASK2METRIC[task],
+        mode="max",
+        max_t=train_args.get("epochs") or DEFAULT_CFG_DICT["epochs"] or 100,
+        grace_period=grace_period,
+        reduction_factor=3,
+    )
+
+    # Define the callbacks for the hyperparameter search
+    tuner_callbacks = [WandbLoggerCallback(project="YOLOv8-tune")] if wandb else []
+
+    # Create the Ray Tune hyperparameter search tuner
+    tune_dir = get_save_dir(
+        get_cfg(
+            DEFAULT_CFG,
+            {**train_args, **{"exist_ok": train_args.pop("resume", False)}},  # resume w/ same tune_dir
+        ),
+        name=train_args.pop("name", "tune"),  # runs/{task}/{tune_dir}
+    )  # must be absolute dir
+    tune_dir.mkdir(parents=True, exist_ok=True)
+    if tune.Tuner.can_restore(tune_dir):
+        LOGGER.info(f"{colorstr('Tuner: ')} Resuming tuning run {tune_dir}...")
+        tuner = tune.Tuner.restore(str(tune_dir), trainable=trainable_with_resources, resume_errored=True)
+    else:
+        tuner = tune.Tuner(
+            trainable_with_resources,
+            param_space=space,
+            tune_config=tune.TuneConfig(
+                scheduler=asha_scheduler,
+                num_samples=max_samples,
+                trial_name_creator=lambda trial: f"{trial.trainable_name}_{trial.trial_id}",
+                trial_dirname_creator=lambda trial: f"{trial.trainable_name}_{trial.trial_id}",
+            ),
+            run_config=RunConfig(callbacks=tuner_callbacks, storage_path=tune_dir.parent, name=tune_dir.name),
+        )
+
+    # Run the hyperparameter search
+    tuner.fit()
+
+    # Get the results of the hyperparameter search
+    results = tuner.get_results()
+
+    # Shut down Ray to clean up workers
+    ray.shutdown()
+
+    return results

ultralytics/utils/uploads.py

@@ -0,0 +1,115 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+"""Upload utilities for Ultralytics, mirroring downloads.py patterns."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from time import sleep
+
+from ultralytics.utils import LOGGER, TQDM
+
+
+class _ProgressReader:
+    """File wrapper that reports read progress for upload monitoring."""
+
+    def __init__(self, file_path, pbar):
+        self.file = open(file_path, "rb")
+        self.pbar = pbar
+        self._size = os.path.getsize(file_path)
+
+    def read(self, size=-1):
+        """Read data and update progress bar."""
+        data = self.file.read(size)
+        if data and self.pbar:
+            self.pbar.update(len(data))
+        return data
+
+    def __len__(self):
+        """Return file size for Content-Length header."""
+        return self._size
+
+    def close(self):
+        """Close the file."""
+        self.file.close()
+
+
+def safe_upload(
+    file: str | Path,
+    url: str,
+    headers: dict | None = None,
+    retry: int = 2,
+    timeout: int = 600,
+    progress: bool = False,
+) -> bool:
+    """Upload a file to a URL with retry logic and optional progress bar.
+
+    Args:
+        file (str | Path): Path to the file to upload.
+        url (str): The URL endpoint to upload the file to (e.g., signed GCS URL).
+        headers (dict, optional): Additional headers to include in the request.
+        retry (int, optional): Number of retry attempts on failure (default: 2 for 3 total attempts).
+        timeout (int, optional): Request timeout in seconds.
+        progress (bool, optional): Whether to display a progress bar during upload.
+
+    Returns:
+        (bool): True if upload succeeded, False otherwise.
+
+    Examples:
+        >>> from ultralytics.utils.uploads import safe_upload
+        >>> success = safe_upload("model.pt", "https://storage.googleapis.com/...", progress=True)
+    """
+    import requests
+
+    file = Path(file)
+    if not file.exists():
+        raise FileNotFoundError(f"File not found: {file}")
+
+    file_size = file.stat().st_size
+    desc = f"Uploading {file.name}"
+
+    # Prepare headers (Content-Length set automatically from file size)
+    upload_headers = {"Content-Type": "application/octet-stream"}
+    if headers:
+        upload_headers.update(headers)
+
+    last_error = None
+    for attempt in range(retry + 1):
+        pbar = None
+        reader = None
+        try:
+            if progress:
+                pbar = TQDM(total=file_size, desc=desc, unit="B", unit_scale=True, unit_divisor=1024)
+            reader = _ProgressReader(file, pbar)
+
+            r = requests.put(url, data=reader, headers=upload_headers, timeout=timeout)
+            r.raise_for_status()
+            reader.close()
+            reader = None  # Prevent double-close in finally
+            if pbar:
+                pbar.close()
+                pbar = None
+            LOGGER.info(f"Uploaded {file.name} ✅")
+            return True
+
+        except requests.exceptions.HTTPError as e:
+            status = e.response.status_code if e.response is not None else 0
+            if 400 <= status < 500 and status not in {408, 429}:
+                LOGGER.warning(f"{desc} failed: {status} {getattr(e.response, 'reason', '')}")
+                return False
+            last_error = f"HTTP {status}"
+        except Exception as e:
+            last_error = str(e)
+        finally:
+            if reader:
+                reader.close()
+            if pbar:
+                pbar.close()
+
+        if attempt < retry:
+            wait_time = 2 ** (attempt + 1)
+            LOGGER.warning(f"{desc} failed ({last_error}), retrying {attempt + 1}/{retry} in {wait_time}s...")
+            sleep(wait_time)
+
+    LOGGER.warning(f"{desc} failed after {retry + 1} attempts: {last_error}")
+    return False