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

ultralytics/utils/downloads.py

@@ -43,8 +43,7 @@ GITHUB_ASSETS_STEMS = frozenset(k.rpartition(".")[0] for k in GITHUB_ASSETS_NAME
 
 
 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.
+    """Validate if the given string is a URL and optionally check if the URL exists online.
 
     Args:
         url (str): The string to be validated as a URL.
@@ -71,8 +70,7 @@ def is_url(url: str | Path, check: bool = False) -> bool:
 
 
 def delete_dsstore(path: str | Path, files_to_delete: tuple[str, ...] = (".DS_Store", "__MACOSX")) -> None:
-    """
-    Delete all specified system files in a directory.
+    """Delete all specified system files in a directory.
 
     Args:
         path (str | Path): The directory path where the files should be deleted.
@@ -83,7 +81,7 @@ def delete_dsstore(path: str | Path, files_to_delete: tuple[str, ...] = (".DS_St
         >>> delete_dsstore("path/to/dir")
 
     Notes:
-        ".DS_store" files are created by the Apple operating system and contain metadata about folders and files. They
+        ".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:
@@ -99,8 +97,7 @@ def zip_directory(
     exclude: tuple[str, ...] = (".DS_Store", "__MACOSX"),
     progress: bool = True,
 ) -> Path:
-    """
-    Zip the contents of a directory, excluding specified files.
+    """Zip the contents of a directory, excluding specified files.
 
     The resulting zip file is named after the directory and placed alongside it.
 
@@ -142,12 +139,11 @@ def unzip_file(
     exist_ok: bool = False,
     progress: bool = True,
 ) -> Path:
-    """
-    Unzip a *.zip file to the specified path, excluding specified files.
+    """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.
+    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.
@@ -183,7 +179,7 @@ def unzip_file(
         if unzip_as_dir:
             # Zip has 1 top-level directory
             extract_path = path  # i.e. ../datasets
-            path = Path(path) / list(top_level_dirs)[0]  # i.e. extract coco8/ dir to ../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/
@@ -210,8 +206,7 @@ def check_disk_space(
     sf: float = 1.5,
     hard: bool = True,
 ) -> bool:
-    """
-    Check if there is sufficient disk space to download and store a file.
+    """Check if there is sufficient disk space to download and store a file.
 
     Args:
         file_bytes (int): The file size in bytes.
@@ -222,7 +217,7 @@ def check_disk_space(
     Returns:
         (bool): True if there is sufficient disk space, False otherwise.
     """
-    total, used, free = shutil.disk_usage(path)  # bytes
+    _total, _used, free = shutil.disk_usage(path)  # bytes
     if file_bytes * sf < free:
         return True  # sufficient space
 
@@ -238,8 +233,7 @@ def check_disk_space(
 
 
 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.
+    """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.
@@ -289,16 +283,15 @@ def safe_download(
     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
+    """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): The URL of the file to be downloaded.
-        file (str, 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.
+        file (str, 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.
@@ -397,8 +390,7 @@ def get_github_assets(
     version: str = "latest",
     retry: bool = False,
 ) -> tuple[str, list[str]]:
-    """
-    Retrieve the specified version's tag and assets from a GitHub repository.
+    """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.
 
@@ -435,8 +427,7 @@ def attempt_download_asset(
     release: str = "v8.3.0",
     **kwargs,
 ) -> str:
-    """
-    Attempt to download a file from GitHub release assets if it is not found locally.
+    """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.
@@ -495,8 +486,7 @@ def download(
     retry: int = 3,
     exist_ok: bool = False,
 ) -> None:
-    """
-    Download files from specified URLs to a given directory.
+    """Download files from specified URLs to a given directory.
 
     Supports concurrent downloads if multiple threads are specified.
 

ultralytics/utils/errors.py

@@ -4,11 +4,10 @@ from ultralytics.utils import emojis
 
 
 class HUBModelError(Exception):
-    """
-    Exception raised when a model cannot be found or retrieved from Ultralytics HUB.
+    """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.
+    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.
@@ -25,19 +24,12 @@ class HUBModelError(Exception):
     """
 
     def __init__(self, message: str = "Model not found. Please check model URL and try again."):
-        """
-        Initialize a HUBModelError exception.
+        """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.
+        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.
-
-        Examples:
-            >>> try:
-            ...     raise HUBModelError("Custom model error message")
-            ... except HUBModelError as e:
-            ...     print(e)
         """
         super().__init__(emojis(message))

ultralytics/utils/events.py

@@ -24,8 +24,7 @@ def _post(url: str, data: dict, timeout: float = 5.0) -> None:
 
 
 class Events:
-    """
-    Collect and send anonymous usage analytics with rate-limiting.
+    """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
@@ -71,8 +70,7 @@ class Events:
         )
 
     def __call__(self, cfg, device=None) -> None:
-        """
-        Queue an event and flush the queue asynchronously when the rate limit elapses.
+        """Queue an event and flush the queue asynchronously when the rate limit elapses.
 
         Args:
             cfg (IterableSimpleNamespace): The configuration object containing mode and task information.

ultralytics/utils/export/__init__.py

@@ -1,239 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from __future__ import annotations
+from .engine import onnx2engine, torch2onnx
+from .imx import torch2imx
+from .tensorflow import keras2pb, onnx2saved_model, pb2tfjs, tflite2edgetpu
 
-import json
-from pathlib import Path
-
-import torch
-
-from ultralytics.utils import IS_JETSON, LOGGER
-
-from .imx import torch2imx  # noqa
-
-
-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.
-    """
-    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,
-    )
-
-
-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, optional): Path to save the generated TensorRT engine file.
-        workspace (int, optional): 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, optional): DLA core to use (Jetson devices only).
-        dataset (ultralytics.data.build.InfiniteDataLoader, optional): Dataset for INT8 calibration.
-        metadata (dict, optional): 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  # noqa
-
-    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:
-            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
-    build = builder.build_serialized_network if is_trt10 else builder.build_engine
-    with build(network, config) as engine, open(engine_file, "wb") as t:
-        # Metadata
-        if metadata is not None:
-            meta = json.dumps(metadata)
-            t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
-            t.write(meta.encode())
-        # Model
-        t.write(engine if is_trt10 else engine.serialize())
+__all__ = ["keras2pb", "onnx2engine", "onnx2saved_model", "pb2tfjs", "tflite2edgetpu", "torch2imx", "torch2onnx"]

ultralytics/utils/export/engine.py

@@ -0,0 +1,237 @@
+# 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, optional): Path to save the generated TensorRT engine file.
+        workspace (int, optional): 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, optional): DLA core to use (Jetson devices only).
+        dataset (ultralytics.data.build.InfiniteDataLoader, optional): Dataset for INT8 calibration.
+        metadata (dict, optional): 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:
+            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
+    build = builder.build_serialized_network if is_trt10 else builder.build_engine
+    with build(network, config) as engine, open(engine_file, "wb") as t:
+        # Metadata
+        if metadata is not None:
+            meta = json.dumps(metadata)
+            t.write(len(meta).to_bytes(4, byteorder="little", signed=True))
+            t.write(meta.encode())
+        # Model
+        t.write(engine if is_trt10 else engine.serialize())