dgenerate-ultralytics-headless 8.3.134__py3-none-any.whl

ultralytics/utils/export.py

@@ -0,0 +1,219 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import json
+from pathlib import Path
+
+import torch
+
+from ultralytics.utils import IS_JETSON, LOGGER
+
+
+def export_onnx(
+    torch_model,
+    im,
+    onnx_file,
+    opset=14,
+    input_names=["images"],
+    output_names=["output0"],
+    dynamic=False,
+):
+    """
+    Exports 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): List of input tensor names.
+        output_names (list): List of output tensor names.
+        dynamic (bool | dict, optional): Whether to enable dynamic axes. Defaults to False.
+
+    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 export_engine(
+    onnx_file,
+    engine_file=None,
+    workspace=None,
+    half=False,
+    int8=False,
+    dynamic=False,
+    shape=(1, 3, 640, 640),
+    dla=None,
+    dataset=None,
+    metadata=None,
+    verbose=False,
+    prefix="",
+):
+    """
+    Exports 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. Defaults to None.
+        half (bool, optional): Enable FP16 precision. Defaults to False.
+        int8 (bool, optional): Enable INT8 precision. Defaults to False.
+        dynamic (bool, optional): Enable dynamic input shapes. Defaults to False.
+        shape (tuple, optional): Input shape (batch, channels, height, width). Defaults to (1, 3, 640, 640).
+        dla (int, optional): DLA core to use (Jetson devices only). Defaults to None.
+        dataset (ultralytics.data.build.InfiniteDataLoader, optional): Dataset for INT8 calibration. Defaults to None.
+        metadata (dict, optional): Metadata to include in the engine file. Defaults to None.
+        verbose (bool, optional): Enable verbose logging. Defaults to False.
+        prefix (str, optional): Prefix for log messages. Defaults to "".
+
+    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 = int((workspace or 0) * (1 << 30))
+    is_trt10 = int(trt.__version__.split(".")[0]) >= 10  # is TensorRT >= 10
+    if is_trt10 and workspace > 0:
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace)
+    elif workspace > 0:  # TensorRT versions 7, 8
+        config.max_workspace_size = workspace
+    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:
+        if shape[0] <= 1:
+            LOGGER.warning(f"{prefix} 'dynamic=True' model requires max batch size, i.e. 'batch=16'")
+        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)
+
+    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.set_calibration_profile(profile)
+        config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
+
+        class EngineCalibrator(trt.IInt8Calibrator):
+            """
+            Custom INT8 calibrator for TensorRT.
+
+            Args:
+                dataset (object): Dataset for calibration.
+                batch (int): Batch size for calibration.
+                cache (str, optional): Path to save the calibration cache. Defaults to "".
+            """
+
+            def __init__(
+                self,
+                dataset,  # ultralytics.data.build.InfiniteDataLoader
+                cache: str = "",
+            ) -> None:
+                trt.IInt8Calibrator.__init__(self)
+                self.dataset = dataset
+                self.data_iter = iter(dataset)
+                self.algo = 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:
+                """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 [] or None, signal to TensorRT there is no calibration data remaining
+                    return None
+
+            def read_calibration_cache(self) -> bytes:
+                """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) -> 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())

ultralytics/utils/files.py

@@ -0,0 +1,221 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+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):
+        """Sets the working directory to 'new_dir' upon instantiation for use with context managers or decorators."""
+        self.dir = new_dir  # new dir
+        self.cwd = Path.cwd().resolve()  # current dir
+
+    def __enter__(self):
+        """Changes 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):  # noqa
+        """Restores the original working directory when exiting the context."""
+        os.chdir(self.cwd)
+
+
+@contextmanager
+def spaces_in_path(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 spaces replaced by underscores if spaces were present, otherwise the original path.
+
+    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():
+                # tmp_path.mkdir(parents=True, exist_ok=True)
+                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, exist_ok=False, sep="", mkdir=False):
+    """
+    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): If True, the path will not be incremented and returned as-is.
+        sep (str): Separator to use between the path and the incrementation number.
+        mkdir (bool): 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=__file__):
+    """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=__file__):
+    """Returns 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):
+    """Returns the size of a file or directory in megabytes (MB)."""
+    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="."):
+    """Returns 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=("yolo11n.pt",), source_dir=Path("."), update_names=False):
+    """
+    Update and re-save specified YOLO models in an 'updated_models' subdirectory.
+
+    Args:
+        model_names (Tuple[str, ...]): Model filenames to update.
+        source_dir (Path): Directory containing models and target subdirectory.
+        update_names (bool): 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 = ("yolo11n.pt", "yolov8s.pt")
+        >>> update_models(model_names, source_dir=Path("/models"), update_names=True)
+    """
+    from ultralytics import YOLO
+    from ultralytics.nn.autobackend import default_class_names
+
+    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
+        print(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()
+        print(f"Re-saving {model_name} model to {save_path}")
+        model.save(save_path)