bplusplus 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

bplusplus/yolov5detect/export.py

@@ -1,1530 +0,0 @@
-# Ultralytics YOLOv5 🚀, AGPL-3.0 license
-"""
-Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit.
-
-Format                      | `export.py --include`         | Model
----                         | ---                           | ---
-PyTorch                     | -                             | yolov5s.pt
-TorchScript                 | `torchscript`                 | yolov5s.torchscript
-ONNX                        | `onnx`                        | yolov5s.onnx
-OpenVINO                    | `openvino`                    | yolov5s_openvino_model/
-TensorRT                    | `engine`                      | yolov5s.engine
-CoreML                      | `coreml`                      | yolov5s.mlmodel
-TensorFlow SavedModel       | `saved_model`                 | yolov5s_saved_model/
-TensorFlow GraphDef         | `pb`                          | yolov5s.pb
-TensorFlow Lite             | `tflite`                      | yolov5s.tflite
-TensorFlow Edge TPU         | `edgetpu`                     | yolov5s_edgetpu.tflite
-TensorFlow.js               | `tfjs`                        | yolov5s_web_model/
-PaddlePaddle                | `paddle`                      | yolov5s_paddle_model/
-
-Requirements:
-    $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu  # CPU
-    $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow  # GPU
-
-Usage:
-    $ python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
-
-Inference:
-    $ python detect.py --weights yolov5s.pt                 # PyTorch
-                                 yolov5s.torchscript        # TorchScript
-                                 yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
-                                 yolov5s_openvino_model     # OpenVINO
-                                 yolov5s.engine             # TensorRT
-                                 yolov5s.mlmodel            # CoreML (macOS-only)
-                                 yolov5s_saved_model        # TensorFlow SavedModel
-                                 yolov5s.pb                 # TensorFlow GraphDef
-                                 yolov5s.tflite             # TensorFlow Lite
-                                 yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
-                                 yolov5s_paddle_model       # PaddlePaddle
-
-TensorFlow.js:
-    $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example
-    $ npm install
-    $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model
-    $ npm start
-"""
-
-import argparse
-import contextlib
-import json
-import os
-import platform
-import re
-import subprocess
-import sys
-import time
-import warnings
-from pathlib import Path
-
-import pandas as pd
-import torch
-from torch.utils.mobile_optimizer import optimize_for_mobile
-
-FILE = Path(__file__).resolve()
-ROOT = FILE.parents[0]  # YOLOv5 root directory
-if str(ROOT) not in sys.path:
-    sys.path.append(str(ROOT))  # add ROOT to PATH
-if platform.system() != "Windows":
-    ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
-
-from models.experimental import attempt_load
-from models.yolo import ClassificationModel, Detect, DetectionModel, SegmentationModel
-from utils.dataloaders import LoadImages
-from utils.general import (
-    LOGGER,
-    Profile,
-    check_dataset,
-    check_img_size,
-    check_requirements,
-    check_version,
-    check_yaml,
-    colorstr,
-    file_size,
-    get_default_args,
-    print_args,
-    url2file,
-    yaml_save,
-)
-from utils.torch_utils import select_device, smart_inference_mode
-
-MACOS = platform.system() == "Darwin"  # macOS environment
-
-
-class iOSModel(torch.nn.Module):
-    """An iOS-compatible wrapper for YOLOv5 models that normalizes input images based on their dimensions."""
-
-    def __init__(self, model, im):
-        """
-        Initializes an iOS compatible model with normalization based on image dimensions.
-
-        Args:
-            model (torch.nn.Module): The PyTorch model to be adapted for iOS compatibility.
-            im (torch.Tensor): An input tensor representing a batch of images with shape (B, C, H, W).
-
-        Returns:
-            None: This method does not return any value.
-
-        Notes:
-            This initializer configures normalization based on the input image dimensions, which is critical for
-            ensuring the model's compatibility and proper functionality on iOS devices. The normalization step
-            involves dividing by the image width if the image is square; otherwise, additional conditions might apply.
-        """
-        super().__init__()
-        b, c, h, w = im.shape  # batch, channel, height, width
-        self.model = model
-        self.nc = model.nc  # number of classes
-        if w == h:
-            self.normalize = 1.0 / w
-        else:
-            self.normalize = torch.tensor([1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h])  # broadcast (slower, smaller)
-            # np = model(im)[0].shape[1]  # number of points
-            # self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]).expand(np, 4)  # explicit (faster, larger)
-
-    def forward(self, x):
-        """
-        Run a forward pass on the input tensor, returning class confidences and normalized coordinates.
-
-        Args:
-            x (torch.Tensor): Input tensor containing the image data with shape (batch, channels, height, width).
-
-        Returns:
-            torch.Tensor: Concatenated tensor with normalized coordinates (xywh), confidence scores (conf),
-            and class probabilities (cls), having shape (N, 4 + 1 + C), where N is the number of predictions,
-            and C is the number of classes.
-
-        Examples:
-            ```python
-            model = iOSModel(pretrained_model, input_image)
-            output = model.forward(torch_input_tensor)
-            ```
-        """
-        xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1)
-        return cls * conf, xywh * self.normalize  # confidence (3780, 80), coordinates (3780, 4)
-
-
-def export_formats():
-    r"""
-    Returns a DataFrame of supported YOLOv5 model export formats and their properties.
-
-    Returns:
-        pandas.DataFrame: A DataFrame containing supported export formats and their properties. The DataFrame
-        includes columns for format name, CLI argument suffix, file extension or directory name, and boolean flags
-        indicating if the export format supports training and detection.
-
-    Examples:
-        ```python
-        formats = export_formats()
-        print(f"Supported export formats:\n{formats}")
-        ```
-
-    Notes:
-        The DataFrame contains the following columns:
-        - Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.).
-        - Include Argument: The argument to use with the export script to include this format.
-        - File Suffix: File extension or directory name associated with the format.
-        - Supports Training: Whether the format supports training.
-        - Supports Detection: Whether the format supports detection.
-    """
-    x = [
-        ["PyTorch", "-", ".pt", True, True],
-        ["TorchScript", "torchscript", ".torchscript", True, True],
-        ["ONNX", "onnx", ".onnx", True, True],
-        ["OpenVINO", "openvino", "_openvino_model", True, False],
-        ["TensorRT", "engine", ".engine", False, True],
-        ["CoreML", "coreml", ".mlpackage", True, False],
-        ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True],
-        ["TensorFlow GraphDef", "pb", ".pb", True, True],
-        ["TensorFlow Lite", "tflite", ".tflite", True, False],
-        ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", False, False],
-        ["TensorFlow.js", "tfjs", "_web_model", False, False],
-        ["PaddlePaddle", "paddle", "_paddle_model", True, True],
-    ]
-    return pd.DataFrame(x, columns=["Format", "Argument", "Suffix", "CPU", "GPU"])
-
-
-def try_export(inner_func):
-    """
-    Log success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export.
-
-    Args:
-        inner_func (Callable): The model export function to be wrapped by the decorator.
-
-    Returns:
-        Callable: The wrapped function that logs execution details. When executed, this wrapper function returns either:
-            - Tuple (str | torch.nn.Module): On success — the file path of the exported model and the model instance.
-            - Tuple (None, None): On failure — None values indicating export failure.
-
-    Examples:
-        ```python
-        @try_export
-        def export_onnx(model, filepath):
-            # implementation here
-            pass
-
-        exported_file, exported_model = export_onnx(yolo_model, 'path/to/save/model.onnx')
-        ```
-
-    Notes:
-        For additional requirements and model export formats, refer to the
-        [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics).
-    """
-    inner_args = get_default_args(inner_func)
-
-    def outer_func(*args, **kwargs):
-        """Logs success/failure and execution details of model export functions wrapped with @try_export decorator."""
-        prefix = inner_args["prefix"]
-        try:
-            with Profile() as dt:
-                f, model = inner_func(*args, **kwargs)
-            LOGGER.info(f"{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)")
-            return f, model
-        except Exception as e:
-            LOGGER.info(f"{prefix} export failure ❌ {dt.t:.1f}s: {e}")
-            return None, None
-
-    return outer_func
-
-
-@try_export
-def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")):
-    """
-    Export a YOLOv5 model to the TorchScript format.
-
-    Args:
-        model (torch.nn.Module): The YOLOv5 model to be exported.
-        im (torch.Tensor): Example input tensor to be used for tracing the TorchScript model.
-        file (Path): File path where the exported TorchScript model will be saved.
-        optimize (bool): If True, applies optimizations for mobile deployment.
-        prefix (str): Optional prefix for log messages. Default is 'TorchScript:'.
-
-    Returns:
-        (str | None, torch.jit.ScriptModule | None): A tuple containing the file path of the exported model
-            (as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements
-            of the tuple will be None.
-
-    Notes:
-        - This function uses tracing to create the TorchScript model.
-        - Metadata, including the input image shape, model stride, and class names, is saved in an extra file (`config.txt`)
-          within the TorchScript model package.
-        - For mobile optimization, refer to the PyTorch tutorial: https://pytorch.org/tutorials/recipes/mobile_interpreter.html
-
-    Example:
-        ```python
-        from pathlib import Path
-        import torch
-        from models.experimental import attempt_load
-        from utils.torch_utils import select_device
-
-        # Load model
-        weights = 'yolov5s.pt'
-        device = select_device('')
-        model = attempt_load(weights, device=device)
-
-        # Example input tensor
-        im = torch.zeros(1, 3, 640, 640).to(device)
-
-        # Export model
-        file = Path('yolov5s.torchscript')
-        export_torchscript(model, im, file, optimize=False)
-        ```
-    """
-    LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...")
-    f = file.with_suffix(".torchscript")
-
-    ts = torch.jit.trace(model, im, strict=False)
-    d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
-    extra_files = {"config.txt": json.dumps(d)}  # torch._C.ExtraFilesMap()
-    if optimize:  # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
-        optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
-    else:
-        ts.save(str(f), _extra_files=extra_files)
-    return f, None
-
-
-@try_export
-def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")):
-    """
-    Export a YOLOv5 model to ONNX format with dynamic axes support and optional model simplification.
-
-    Args:
-        model (torch.nn.Module): The YOLOv5 model to be exported.
-        im (torch.Tensor): A sample input tensor for model tracing, usually the shape is (1, 3, height, width).
-        file (pathlib.Path | str): The output file path where the ONNX model will be saved.
-        opset (int): The ONNX opset version to use for export.
-        dynamic (bool): If True, enables dynamic axes for batch, height, and width dimensions.
-        simplify (bool): If True, applies ONNX model simplification for optimization.
-        prefix (str): A prefix string for logging messages, defaults to 'ONNX:'.
-
-    Returns:
-        tuple[pathlib.Path | str, None]: The path to the saved ONNX model file and None (consistent with decorator).
-
-    Raises:
-        ImportError: If required libraries for export (e.g., 'onnx', 'onnx-simplifier') are not installed.
-        AssertionError: If the simplification check fails.
-
-    Notes:
-        The required packages for this function can be installed via:
-        ```
-        pip install onnx onnx-simplifier onnxruntime onnxruntime-gpu
-        ```
-
-    Example:
-        ```python
-        from pathlib import Path
-        import torch
-        from models.experimental import attempt_load
-        from utils.torch_utils import select_device
-
-        # Load model
-        weights = 'yolov5s.pt'
-        device = select_device('')
-        model = attempt_load(weights, map_location=device)
-
-        # Example input tensor
-        im = torch.zeros(1, 3, 640, 640).to(device)
-
-        # Export model
-        file_path = Path('yolov5s.onnx')
-        export_onnx(model, im, file_path, opset=12, dynamic=True, simplify=True)
-        ```
-    """
-    check_requirements("onnx>=1.12.0")
-    import onnx
-
-    LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...")
-    f = str(file.with_suffix(".onnx"))
-
-    output_names = ["output0", "output1"] if isinstance(model, SegmentationModel) else ["output0"]
-    if dynamic:
-        dynamic = {"images": {0: "batch", 2: "height", 3: "width"}}  # shape(1,3,640,640)
-        if isinstance(model, SegmentationModel):
-            dynamic["output0"] = {0: "batch", 1: "anchors"}  # shape(1,25200,85)
-            dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"}  # shape(1,32,160,160)
-        elif isinstance(model, DetectionModel):
-            dynamic["output0"] = {0: "batch", 1: "anchors"}  # shape(1,25200,85)
-
-    torch.onnx.export(
-        model.cpu() if dynamic else model,  # --dynamic only compatible with cpu
-        im.cpu() if dynamic else im,
-        f,
-        verbose=False,
-        opset_version=opset,
-        do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
-        input_names=["images"],
-        output_names=output_names,
-        dynamic_axes=dynamic or None,
-    )
-
-    # Checks
-    model_onnx = onnx.load(f)  # load onnx model
-    onnx.checker.check_model(model_onnx)  # check onnx model
-
-    # Metadata
-    d = {"stride": int(max(model.stride)), "names": model.names}
-    for k, v in d.items():
-        meta = model_onnx.metadata_props.add()
-        meta.key, meta.value = k, str(v)
-    onnx.save(model_onnx, f)
-
-    # Simplify
-    if simplify:
-        try:
-            cuda = torch.cuda.is_available()
-            check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnxslim"))
-            import onnxslim
-
-            LOGGER.info(f"{prefix} slimming with onnxslim {onnxslim.__version__}...")
-            model_onnx = onnxslim.slim(model_onnx)
-            onnx.save(model_onnx, f)
-        except Exception as e:
-            LOGGER.info(f"{prefix} simplifier failure: {e}")
-    return f, model_onnx
-
-
-@try_export
-def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")):
-    """
-    Export a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization.
-
-    Args:
-        file (Path): Path to the output file where the OpenVINO model will be saved.
-        metadata (dict): Dictionary including model metadata such as names and strides.
-        half (bool): If True, export the model with FP16 precision.
-        int8 (bool): If True, export the model with INT8 quantization.
-        data (str): Path to the dataset YAML file required for INT8 quantization.
-        prefix (str): Prefix string for logging purposes (default is "OpenVINO:").
-
-    Returns:
-        (str, openvino.runtime.Model | None): The OpenVINO model file path and openvino.runtime.Model object if export is
-            successful; otherwise, None.
-
-    Notes:
-        - Requires `openvino-dev` package version 2023.0 or higher. Install with:
-          `$ pip install openvino-dev>=2023.0`
-        - For INT8 quantization, also requires `nncf` library version 2.5.0 or higher. Install with:
-          `$ pip install nncf>=2.5.0`
-
-    Examples:
-        ```python
-        from pathlib import Path
-        from ultralytics import YOLOv5
-
-        model = YOLOv5('yolov5s.pt')
-        export_openvino(Path('yolov5s.onnx'), metadata={'names': model.names, 'stride': model.stride}, half=True,
-                        int8=False, data='data.yaml')
-        ```
-
-        This will export the YOLOv5 model to OpenVINO with FP16 precision but without INT8 quantization, saving it to
-        the specified file path.
-    """
-    check_requirements("openvino-dev>=2023.0")  # requires openvino-dev: https://pypi.org/project/openvino-dev/
-    import openvino.runtime as ov  # noqa
-    from openvino.tools import mo  # noqa
-
-    LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
-    f = str(file).replace(file.suffix, f"_{'int8_' if int8 else ''}openvino_model{os.sep}")
-    f_onnx = file.with_suffix(".onnx")
-    f_ov = str(Path(f) / file.with_suffix(".xml").name)
-
-    ov_model = mo.convert_model(f_onnx, model_name=file.stem, framework="onnx", compress_to_fp16=half)  # export
-
-    if int8:
-        check_requirements("nncf>=2.5.0")  # requires at least version 2.5.0 to use the post-training quantization
-        import nncf
-        import numpy as np
-
-        from utils.dataloaders import create_dataloader
-
-        def gen_dataloader(yaml_path, task="train", imgsz=640, workers=4):
-            """Generates a DataLoader for model training or validation based on the given YAML dataset configuration."""
-            data_yaml = check_yaml(yaml_path)
-            data = check_dataset(data_yaml)
-            dataloader = create_dataloader(
-                data[task], imgsz=imgsz, batch_size=1, stride=32, pad=0.5, single_cls=False, rect=False, workers=workers
-            )[0]
-            return dataloader
-
-        # noqa: F811
-
-        def transform_fn(data_item):
-            """
-            Quantization transform function.
-
-            Extracts and preprocess input data from dataloader item for quantization.
-
-            Args:
-               data_item: Tuple with data item produced by DataLoader during iteration
-
-            Returns:
-                input_tensor: Input data for quantization
-            """
-            assert data_item[0].dtype == torch.uint8, "input image must be uint8 for the quantization preprocessing"
-
-            img = data_item[0].numpy().astype(np.float32)  # uint8 to fp16/32
-            img /= 255.0  # 0 - 255 to 0.0 - 1.0
-            return np.expand_dims(img, 0) if img.ndim == 3 else img
-
-        ds = gen_dataloader(data)
-        quantization_dataset = nncf.Dataset(ds, transform_fn)
-        ov_model = nncf.quantize(ov_model, quantization_dataset, preset=nncf.QuantizationPreset.MIXED)
-
-    ov.serialize(ov_model, f_ov)  # save
-    yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata)  # add metadata.yaml
-    return f, None
-
-
-@try_export
-def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")):
-    """
-    Export a YOLOv5 PyTorch model to PaddlePaddle format using X2Paddle, saving the converted model and metadata.
-
-    Args:
-        model (torch.nn.Module): The YOLOv5 model to be exported.
-        im (torch.Tensor): Input tensor used for model tracing during export.
-        file (pathlib.Path): Path to the source file to be converted.
-        metadata (dict): Additional metadata to be saved alongside the model.
-        prefix (str): Prefix for logging information.
-
-    Returns:
-        tuple (str, None): A tuple where the first element is the path to the saved PaddlePaddle model, and the
-        second element is None.
-
-    Examples:
-        ```python
-        from pathlib import Path
-        import torch
-
-        # Assume 'model' is a pre-trained YOLOv5 model and 'im' is an example input tensor
-        model = ...  # Load your model here
-        im = torch.randn((1, 3, 640, 640))  # Dummy input tensor for tracing
-        file = Path("yolov5s.pt")
-        metadata = {"stride": 32, "names": ["person", "bicycle", "car", "motorbike"]}
-
-        export_paddle(model=model, im=im, file=file, metadata=metadata)
-        ```
-
-    Notes:
-        Ensure that `paddlepaddle` and `x2paddle` are installed, as these are required for the export function. You can
-        install them via pip:
-        ```
-        $ pip install paddlepaddle x2paddle
-        ```
-    """
-    check_requirements(("paddlepaddle", "x2paddle"))
-    import x2paddle
-    from x2paddle.convert import pytorch2paddle
-
-    LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
-    f = str(file).replace(".pt", f"_paddle_model{os.sep}")
-
-    pytorch2paddle(module=model, save_dir=f, jit_type="trace", input_examples=[im])  # export
-    yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata)  # add metadata.yaml
-    return f, None
-
-
-@try_export
-def export_coreml(model, im, file, int8, half, nms, mlmodel, prefix=colorstr("CoreML:")):
-    """
-    Export a YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support.
-
-    Args:
-        model (torch.nn.Module): The YOLOv5 model to be exported.
-        im (torch.Tensor): Example input tensor to trace the model.
-        file (pathlib.Path): Path object where the CoreML model will be saved.
-        int8 (bool): Flag indicating whether to use INT8 quantization (default is False).
-        half (bool): Flag indicating whether to use FP16 quantization (default is False).
-        nms (bool): Flag indicating whether to include Non-Maximum Suppression (default is False).
-        mlmodel (bool): Flag indicating whether to export as older *.mlmodel format (default is False).
-        prefix (str): Prefix string for logging purposes (default is 'CoreML:').
-
-    Returns:
-        tuple[pathlib.Path | None, None]: The path to the saved CoreML model file, or (None, None) if there is an error.
-
-    Notes:
-        The exported CoreML model will be saved with a .mlmodel extension.
-        Quantization is supported only on macOS.
-
-    Example:
-        ```python
-        from pathlib import Path
-        import torch
-        from models.yolo import Model
-        model = Model(cfg, ch=3, nc=80)
-        im = torch.randn(1, 3, 640, 640)
-        file = Path("yolov5s_coreml")
-        export_coreml(model, im, file, int8=False, half=False, nms=True, mlmodel=False)
-        ```
-    """
-    check_requirements("coremltools")
-    import coremltools as ct
-
-    LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
-    if mlmodel:
-        f = file.with_suffix(".mlmodel")
-        convert_to = "neuralnetwork"
-        precision = None
-    else:
-        f = file.with_suffix(".mlpackage")
-        convert_to = "mlprogram"
-        precision = ct.precision.FLOAT16 if half else ct.precision.FLOAT32
-    if nms:
-        model = iOSModel(model, im)
-    ts = torch.jit.trace(model, im, strict=False)  # TorchScript model
-    ct_model = ct.convert(
-        ts,
-        inputs=[ct.ImageType("image", shape=im.shape, scale=1 / 255, bias=[0, 0, 0])],
-        convert_to=convert_to,
-        compute_precision=precision,
-    )
-    bits, mode = (8, "kmeans") if int8 else (16, "linear") if half else (32, None)
-    if bits < 32:
-        if mlmodel:
-            with warnings.catch_warnings():
-                warnings.filterwarnings(
-                    "ignore", category=DeprecationWarning
-                )  # suppress numpy==1.20 float warning, fixed in coremltools==7.0
-                ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
-        elif bits == 8:
-            op_config = ct.optimize.coreml.OpPalettizerConfig(mode=mode, nbits=bits, weight_threshold=512)
-            config = ct.optimize.coreml.OptimizationConfig(global_config=op_config)
-            ct_model = ct.optimize.coreml.palettize_weights(ct_model, config)
-    ct_model.save(f)
-    return f, ct_model
-
-
-@try_export
-def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr("TensorRT:")):
-    """
-    Export a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0.
-
-    Args:
-        model (torch.nn.Module): YOLOv5 model to be exported.
-        im (torch.Tensor): Input tensor of shape (B, C, H, W).
-        file (pathlib.Path): Path to save the exported model.
-        half (bool): Set to True to export with FP16 precision.
-        dynamic (bool): Set to True to enable dynamic input shapes.
-        simplify (bool): Set to True to simplify the model during export.
-        workspace (int): Workspace size in GB (default is 4).
-        verbose (bool): Set to True for verbose logging output.
-        prefix (str): Log message prefix.
-
-    Returns:
-        (pathlib.Path, None): Tuple containing the path to the exported model and None.
-
-    Raises:
-        AssertionError: If executed on CPU instead of GPU.
-        RuntimeError: If there is a failure in parsing the ONNX file.
-
-    Example:
-        ```python
-        from ultralytics import YOLOv5
-        import torch
-        from pathlib import Path
-
-        model = YOLOv5('yolov5s.pt')  # Load a pre-trained YOLOv5 model
-        input_tensor = torch.randn(1, 3, 640, 640).cuda()  # example input tensor on GPU
-        export_path = Path('yolov5s.engine')  # export destination
-
-        export_engine(model.model, input_tensor, export_path, half=True, dynamic=True, simplify=True, workspace=8, verbose=True)
-        ```
-    """
-    assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`"
-    try:
-        import tensorrt as trt
-    except Exception:
-        if platform.system() == "Linux":
-            check_requirements("nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com")
-        import tensorrt as trt
-
-    if trt.__version__[0] == "7":  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
-        grid = model.model[-1].anchor_grid
-        model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid]
-        export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
-        model.model[-1].anchor_grid = grid
-    else:  # TensorRT >= 8
-        check_version(trt.__version__, "8.0.0", hard=True)  # require tensorrt>=8.0.0
-        export_onnx(model, im, file, 12, dynamic, simplify)  # opset 12
-    onnx = file.with_suffix(".onnx")
-
-    LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")
-    is_trt10 = int(trt.__version__.split(".")[0]) >= 10  # is TensorRT >= 10
-    assert onnx.exists(), f"failed to export ONNX file: {onnx}"
-    f = file.with_suffix(".engine")  # TensorRT engine file
-    logger = trt.Logger(trt.Logger.INFO)
-    if verbose:
-        logger.min_severity = trt.Logger.Severity.VERBOSE
-
-    builder = trt.Builder(logger)
-    config = builder.create_builder_config()
-    if is_trt10:
-        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30)
-    else:  # TensorRT versions 7, 8
-        config.max_workspace_size = workspace * 1 << 30
-    flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
-    network = builder.create_network(flag)
-    parser = trt.OnnxParser(network, logger)
-    if not parser.parse_from_file(str(onnx)):
-        raise RuntimeError(f"failed to load ONNX file: {onnx}")
-
-    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 im.shape[0] <= 1:
-            LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument")
-        profile = builder.create_optimization_profile()
-        for inp in inputs:
-            profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape)
-        config.add_optimization_profile(profile)
-
-    LOGGER.info(f"{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}")
-    if builder.platform_has_fast_fp16 and half:
-        config.set_flag(trt.BuilderFlag.FP16)
-
-    build = builder.build_serialized_network if is_trt10 else builder.build_engine
-    with build(network, config) as engine, open(f, "wb") as t:
-        t.write(engine if is_trt10 else engine.serialize())
-    return f, None
-
-
-@try_export
-def export_saved_model(
-    model,
-    im,
-    file,
-    dynamic,
-    tf_nms=False,
-    agnostic_nms=False,
-    topk_per_class=100,
-    topk_all=100,
-    iou_thres=0.45,
-    conf_thres=0.25,
-    keras=False,
-    prefix=colorstr("TensorFlow SavedModel:"),
-):
-    """
-    Export a YOLOv5 model to the TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression
-    (NMS).
-
-    Args:
-        model (torch.nn.Module): The PyTorch model to convert.
-        im (torch.Tensor): Sample input tensor with shape (B, C, H, W) for tracing.
-        file (pathlib.Path): File path to save the exported model.
-        dynamic (bool): Flag to indicate whether dynamic axes should be used.
-        tf_nms (bool, optional): Enable TensorFlow non-maximum suppression (NMS). Default is False.
-        agnostic_nms (bool, optional): Enable class-agnostic NMS. Default is False.
-        topk_per_class (int, optional): Top K detections per class to keep before applying NMS. Default is 100.
-        topk_all (int, optional): Top K detections across all classes to keep before applying NMS. Default is 100.
-        iou_thres (float, optional): IoU threshold for NMS. Default is 0.45.
-        conf_thres (float, optional): Confidence threshold for detections. Default is 0.25.
-        keras (bool, optional): Save the model in Keras format if True. Default is False.
-        prefix (str, optional): Prefix for logging messages. Default is "TensorFlow SavedModel:".
-
-    Returns:
-        tuple[str, tf.keras.Model | None]: A tuple containing the path to the saved model folder and the Keras model instance,
-        or None if TensorFlow export fails.
-
-    Notes:
-        - The method supports TensorFlow versions up to 2.15.1.
-        - TensorFlow NMS may not be supported in older TensorFlow versions.
-        - If the TensorFlow version exceeds 2.13.1, it might cause issues when exporting to TFLite.
-          Refer to: https://github.com/ultralytics/yolov5/issues/12489
-
-    Example:
-        ```python
-        model, im = ...  # Initialize your PyTorch model and input tensor
-        export_saved_model(model, im, Path("yolov5_saved_model"), dynamic=True)
-        ```
-    """
-    # YOLOv5 TensorFlow SavedModel export
-    try:
-        import tensorflow as tf
-    except Exception:
-        check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}<=2.15.1")
-
-        import tensorflow as tf
-    from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
-
-    from models.tf import TFModel
-
-    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
-    if tf.__version__ > "2.13.1":
-        helper_url = "https://github.com/ultralytics/yolov5/issues/12489"
-        LOGGER.info(
-            f"WARNING ⚠️ using Tensorflow {tf.__version__} > 2.13.1 might cause issue when exporting the model to tflite {helper_url}"
-        )  # handling issue https://github.com/ultralytics/yolov5/issues/12489
-    f = str(file).replace(".pt", "_saved_model")
-    batch_size, ch, *imgsz = list(im.shape)  # BCHW
-
-    tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz)
-    im = tf.zeros((batch_size, *imgsz, ch))  # BHWC order for TensorFlow
-    _ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
-    inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size)
-    outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)
-    keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)
-    keras_model.trainable = False
-    keras_model.summary()
-    if keras:
-        keras_model.save(f, save_format="tf")
-    else:
-        spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)
-        m = tf.function(lambda x: keras_model(x))  # full model
-        m = m.get_concrete_function(spec)
-        frozen_func = convert_variables_to_constants_v2(m)
-        tfm = tf.Module()
-        tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec])
-        tfm.__call__(im)
-        tf.saved_model.save(
-            tfm,
-            f,
-            options=tf.saved_model.SaveOptions(experimental_custom_gradients=False)
-            if check_version(tf.__version__, "2.6")
-            else tf.saved_model.SaveOptions(),
-        )
-    return f, keras_model
-
-
-@try_export
-def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")):
-    """
-    Export YOLOv5 model to TensorFlow GraphDef (*.pb) format.
-
-    Args:
-        keras_model (tf.keras.Model): The Keras model to be converted.
-        file (Path): The output file path where the GraphDef will be saved.
-        prefix (str): Optional prefix string; defaults to a colored string indicating TensorFlow GraphDef export status.
-
-    Returns:
-        Tuple[Path, None]: The file path where the GraphDef model was saved and a None placeholder.
-
-    Notes:
-        For more details, refer to the guide on frozen graphs: https://github.com/leimao/Frozen_Graph_TensorFlow
-
-    Example:
-        ```python
-        from pathlib import Path
-        keras_model = ...  # assume an existing Keras model
-        file = Path("model.pb")
-        export_pb(keras_model, file)
-        ```
-    """
-    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__}...")
-    f = file.with_suffix(".pb")
-
-    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(f.parent), name=f.name, as_text=False)
-    return f, None
-
-
-@try_export
-def export_tflite(
-    keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:")
-):
-    # YOLOv5 TensorFlow Lite export
-    """
-    Export a YOLOv5 model to TensorFlow Lite format with optional INT8 quantization and NMS support.
-
-    Args:
-        keras_model (tf.keras.Model): The Keras model to be exported.
-        im (torch.Tensor): An input image tensor for normalization and model tracing.
-        file (Path): The file path to save the TensorFlow Lite model.
-        int8 (bool): Enables INT8 quantization if True.
-        per_tensor (bool): If True, disables per-channel quantization.
-        data (str): Path to the dataset for representative dataset generation in INT8 quantization.
-        nms (bool): Enables Non-Maximum Suppression (NMS) if True.
-        agnostic_nms (bool): Enables class-agnostic NMS if True.
-        prefix (str): Prefix for log messages.
-
-    Returns:
-        (str | None, tflite.Model | None): The file path of the exported TFLite model and the TFLite model instance, or None
-        if the export failed.
-
-    Example:
-        ```python
-        from pathlib import Path
-        import torch
-        import tensorflow as tf
-
-        # Load a Keras model wrapping a YOLOv5 model
-        keras_model = tf.keras.models.load_model('path/to/keras_model.h5')
-
-        # Example input tensor
-        im = torch.zeros(1, 3, 640, 640)
-
-        # Export the model
-        export_tflite(keras_model, im, Path('model.tflite'), int8=True, per_tensor=False, data='data/coco.yaml',
-                      nms=True, agnostic_nms=False)
-        ```
-
-    Notes:
-        - Ensure TensorFlow and TensorFlow Lite dependencies are installed.
-        - INT8 quantization requires a representative dataset to achieve optimal accuracy.
-        - TensorFlow Lite models are suitable for efficient inference on mobile and edge devices.
-    """
-    import tensorflow as tf
-
-    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
-    batch_size, ch, *imgsz = list(im.shape)  # BCHW
-    f = str(file).replace(".pt", "-fp16.tflite")
-
-    converter = tf.lite.TFLiteConverter.from_keras_model(keras_model)
-    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
-    converter.target_spec.supported_types = [tf.float16]
-    converter.optimizations = [tf.lite.Optimize.DEFAULT]
-    if int8:
-        from models.tf import representative_dataset_gen
-
-        dataset = LoadImages(check_dataset(check_yaml(data))["train"], img_size=imgsz, auto=False)
-        converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100)
-        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
-        converter.target_spec.supported_types = []
-        converter.inference_input_type = tf.uint8  # or tf.int8
-        converter.inference_output_type = tf.uint8  # or tf.int8
-        converter.experimental_new_quantizer = True
-        if per_tensor:
-            converter._experimental_disable_per_channel = True
-        f = str(file).replace(".pt", "-int8.tflite")
-    if nms or agnostic_nms:
-        converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
-
-    tflite_model = converter.convert()
-    open(f, "wb").write(tflite_model)
-    return f, None
-
-
-@try_export
-def export_edgetpu(file, prefix=colorstr("Edge TPU:")):
-    """
-    Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler.
-
-    Args:
-        file (Path): Path to the YOLOv5 model file to be exported (.pt format).
-        prefix (str, optional): Prefix for logging messages. Defaults to colorstr("Edge TPU:").
-
-    Returns:
-        tuple[Path, None]: Path to the exported Edge TPU compatible TFLite model, None.
-
-    Raises:
-        AssertionError: If the system is not Linux.
-        subprocess.CalledProcessError: If any subprocess call to install or run the Edge TPU compiler fails.
-
-    Notes:
-        To use this function, ensure you have the Edge TPU compiler installed on your Linux system. You can find
-        installation instructions here: https://coral.ai/docs/edgetpu/compiler/.
-
-    Example:
-        ```python
-        from pathlib import Path
-        file = Path('yolov5s.pt')
-        export_edgetpu(file)
-        ```
-    """
-    cmd = "edgetpu_compiler --version"
-    help_url = "https://coral.ai/docs/edgetpu/compiler/"
-    assert platform.system() == "Linux", f"export only supported on Linux. See {help_url}"
-    if subprocess.run(f"{cmd} > /dev/null 2>&1", shell=True).returncode != 0:
-        LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}")
-        sudo = subprocess.run("sudo --version >/dev/null", shell=True).returncode == 0  # sudo installed on system
-        for c in (
-            "curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -",
-            'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list',
-            "sudo apt-get update",
-            "sudo apt-get install edgetpu-compiler",
-        ):
-            subprocess.run(c if sudo else c.replace("sudo ", ""), shell=True, check=True)
-    ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1]
-
-    LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
-    f = str(file).replace(".pt", "-int8_edgetpu.tflite")  # Edge TPU model
-    f_tfl = str(file).replace(".pt", "-int8.tflite")  # TFLite model
-
-    subprocess.run(
-        [
-            "edgetpu_compiler",
-            "-s",
-            "-d",
-            "-k",
-            "10",
-            "--out_dir",
-            str(file.parent),
-            f_tfl,
-        ],
-        check=True,
-    )
-    return f, None
-
-
-@try_export
-def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
-    """
-    Convert a YOLOv5 model to TensorFlow.js format with optional uint8 quantization.
-
-    Args:
-        file (Path): Path to the YOLOv5 model file to be converted, typically having a ".pt" or ".onnx" extension.
-        int8 (bool): If True, applies uint8 quantization during the conversion process.
-        prefix (str): Optional prefix for logging messages, default is 'TensorFlow.js:' with color formatting.
-
-    Returns:
-        (str, None): Tuple containing the output directory path as a string and None.
-
-    Notes:
-        - This function requires the `tensorflowjs` package. Install it using:
-          ```shell
-          pip install tensorflowjs
-          ```
-        - The converted TensorFlow.js model will be saved in a directory with the "_web_model" suffix appended to the original file name.
-        - The conversion involves running shell commands that invoke the TensorFlow.js converter tool.
-
-    Example:
-        ```python
-        from pathlib import Path
-        file = Path('yolov5.onnx')
-        export_tfjs(file, int8=False)
-        ```
-    """
-    check_requirements("tensorflowjs")
-    import tensorflowjs as tfjs
-
-    LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...")
-    f = str(file).replace(".pt", "_web_model")  # js dir
-    f_pb = file.with_suffix(".pb")  # *.pb path
-    f_json = f"{f}/model.json"  # *.json path
-
-    args = [
-        "tensorflowjs_converter",
-        "--input_format=tf_frozen_model",
-        "--quantize_uint8" if int8 else "",
-        "--output_node_names=Identity,Identity_1,Identity_2,Identity_3",
-        str(f_pb),
-        f,
-    ]
-    subprocess.run([arg for arg in args if arg], check=True)
-
-    json = Path(f_json).read_text()
-    with open(f_json, "w") as j:  # sort JSON Identity_* in ascending order
-        subst = re.sub(
-            r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, '
-            r'"Identity.?.?": {"name": "Identity.?.?"}, '
-            r'"Identity.?.?": {"name": "Identity.?.?"}, '
-            r'"Identity.?.?": {"name": "Identity.?.?"}}}',
-            r'{"outputs": {"Identity": {"name": "Identity"}, '
-            r'"Identity_1": {"name": "Identity_1"}, '
-            r'"Identity_2": {"name": "Identity_2"}, '
-            r'"Identity_3": {"name": "Identity_3"}}}',
-            json,
-        )
-        j.write(subst)
-    return f, None
-
-
-def add_tflite_metadata(file, metadata, num_outputs):
-    """
-    Adds metadata to a TensorFlow Lite (TFLite) model file, supporting multiple outputs according to TensorFlow
-    guidelines.
-
-    Args:
-        file (str): Path to the TFLite model file to which metadata will be added.
-        metadata (dict): Metadata information to be added to the model, structured as required by the TFLite metadata schema.
-            Common keys include "name", "description", "version", "author", and "license".
-        num_outputs (int): Number of output tensors the model has, used to configure the metadata properly.
-
-    Returns:
-        None
-
-    Example:
-        ```python
-        metadata = {
-            "name": "yolov5",
-            "description": "YOLOv5 object detection model",
-            "version": "1.0",
-            "author": "Ultralytics",
-            "license": "Apache License 2.0"
-        }
-        add_tflite_metadata("model.tflite", metadata, num_outputs=4)
-        ```
-
-    Note:
-        TFLite metadata can include information such as model name, version, author, and other relevant details.
-        For more details on the structure of the metadata, refer to TensorFlow Lite
-        [metadata guidelines](https://www.tensorflow.org/lite/models/convert/metadata).
-    """
-    with contextlib.suppress(ImportError):
-        # check_requirements('tflite_support')
-        from tflite_support import flatbuffers
-        from tflite_support import metadata as _metadata
-        from tflite_support import metadata_schema_py_generated as _metadata_fb
-
-        tmp_file = Path("/tmp/meta.txt")
-        with open(tmp_file, "w") as meta_f:
-            meta_f.write(str(metadata))
-
-        model_meta = _metadata_fb.ModelMetadataT()
-        label_file = _metadata_fb.AssociatedFileT()
-        label_file.name = tmp_file.name
-        model_meta.associatedFiles = [label_file]
-
-        subgraph = _metadata_fb.SubGraphMetadataT()
-        subgraph.inputTensorMetadata = [_metadata_fb.TensorMetadataT()]
-        subgraph.outputTensorMetadata = [_metadata_fb.TensorMetadataT()] * num_outputs
-        model_meta.subgraphMetadata = [subgraph]
-
-        b = flatbuffers.Builder(0)
-        b.Finish(model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
-        metadata_buf = b.Output()
-
-        populator = _metadata.MetadataPopulator.with_model_file(file)
-        populator.load_metadata_buffer(metadata_buf)
-        populator.load_associated_files([str(tmp_file)])
-        populator.populate()
-        tmp_file.unlink()
-
-
-def pipeline_coreml(model, im, file, names, y, mlmodel, prefix=colorstr("CoreML Pipeline:")):
-    """
-    Convert a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output
-    shapes, and saving the model.
-
-    Args:
-        model (torch.nn.Module): The YOLOv5 PyTorch model to be converted.
-        im (torch.Tensor): Example input tensor with shape (N, C, H, W), where N is the batch size, C is the number of channels,
-            H is the height, and W is the width.
-        file (Path): Path to save the converted CoreML model.
-        names (dict[int, str]): Dictionary mapping class indices to class names.
-        y (torch.Tensor): Output tensor from the PyTorch model's forward pass.
-        mlmodel (bool): Flag indicating whether to export as older *.mlmodel format (default is False).
-        prefix (str): Custom prefix for logging messages.
-
-    Returns:
-        (Path): Path to the saved CoreML model (.mlmodel).
-
-    Raises:
-        AssertionError: If the number of class names does not match the number of classes in the model.
-
-    Notes:
-        - This function requires `coremltools` to be installed.
-        - Running this function on a non-macOS environment might not support some features.
-        - Flexible input shapes and additional NMS options can be customized within the function.
-
-    Examples:
-        ```python
-        from pathlib import Path
-        import torch
-
-        model = torch.load('yolov5s.pt')  # Load YOLOv5 model
-        im = torch.zeros((1, 3, 640, 640))  # Example input tensor
-
-        names = {0: "person", 1: "bicycle", 2: "car", ...}  # Define class names
-
-        y = model(im)  # Perform forward pass to get model output
-
-        output_file = Path('yolov5s.mlmodel')  # Convert to CoreML
-        pipeline_coreml(model, im, output_file, names, y)
-        ```
-    """
-    import coremltools as ct
-    from PIL import Image
-
-    f = file.with_suffix(".mlmodel") if mlmodel else file.with_suffix(".mlpackage")
-    print(f"{prefix} starting pipeline with coremltools {ct.__version__}...")
-    batch_size, ch, h, w = list(im.shape)  # BCHW
-    t = time.time()
-
-    # YOLOv5 Output shapes
-    spec = model.get_spec()
-    out0, out1 = iter(spec.description.output)
-    if platform.system() == "Darwin":
-        img = Image.new("RGB", (w, h))  # img(192 width, 320 height)
-        # img = torch.zeros((*opt.img_size, 3)).numpy()  # img size(320,192,3) iDetection
-        out = model.predict({"image": img})
-        out0_shape, out1_shape = out[out0.name].shape, out[out1.name].shape
-    else:  # linux and windows can not run model.predict(), get sizes from pytorch output y
-        s = tuple(y[0].shape)
-        out0_shape, out1_shape = (s[1], s[2] - 5), (s[1], 4)  # (3780, 80), (3780, 4)
-
-    # Checks
-    nx, ny = spec.description.input[0].type.imageType.width, spec.description.input[0].type.imageType.height
-    na, nc = out0_shape
-    # na, nc = out0.type.multiArrayType.shape  # number anchors, classes
-    assert len(names) == nc, f"{len(names)} names found for nc={nc}"  # check
-
-    # Define output shapes (missing)
-    out0.type.multiArrayType.shape[:] = out0_shape  # (3780, 80)
-    out1.type.multiArrayType.shape[:] = out1_shape  # (3780, 4)
-    # spec.neuralNetwork.preprocessing[0].featureName = '0'
-
-    # Flexible input shapes
-    # from coremltools.models.neural_network import flexible_shape_utils
-    # s = [] # shapes
-    # s.append(flexible_shape_utils.NeuralNetworkImageSize(320, 192))
-    # s.append(flexible_shape_utils.NeuralNetworkImageSize(640, 384))  # (height, width)
-    # flexible_shape_utils.add_enumerated_image_sizes(spec, feature_name='image', sizes=s)
-    # r = flexible_shape_utils.NeuralNetworkImageSizeRange()  # shape ranges
-    # r.add_height_range((192, 640))
-    # r.add_width_range((192, 640))
-    # flexible_shape_utils.update_image_size_range(spec, feature_name='image', size_range=r)
-
-    # Print
-    print(spec.description)
-
-    # Model from spec
-    weights_dir = None
-    weights_dir = None if mlmodel else str(f / "Data/com.apple.CoreML/weights")
-    model = ct.models.MLModel(spec, weights_dir=weights_dir)
-
-    # 3. Create NMS protobuf
-    nms_spec = ct.proto.Model_pb2.Model()
-    nms_spec.specificationVersion = 5
-    for i in range(2):
-        decoder_output = model._spec.description.output[i].SerializeToString()
-        nms_spec.description.input.add()
-        nms_spec.description.input[i].ParseFromString(decoder_output)
-        nms_spec.description.output.add()
-        nms_spec.description.output[i].ParseFromString(decoder_output)
-
-    nms_spec.description.output[0].name = "confidence"
-    nms_spec.description.output[1].name = "coordinates"
-
-    output_sizes = [nc, 4]
-    for i in range(2):
-        ma_type = nms_spec.description.output[i].type.multiArrayType
-        ma_type.shapeRange.sizeRanges.add()
-        ma_type.shapeRange.sizeRanges[0].lowerBound = 0
-        ma_type.shapeRange.sizeRanges[0].upperBound = -1
-        ma_type.shapeRange.sizeRanges.add()
-        ma_type.shapeRange.sizeRanges[1].lowerBound = output_sizes[i]
-        ma_type.shapeRange.sizeRanges[1].upperBound = output_sizes[i]
-        del ma_type.shape[:]
-
-    nms = nms_spec.nonMaximumSuppression
-    nms.confidenceInputFeatureName = out0.name  # 1x507x80
-    nms.coordinatesInputFeatureName = out1.name  # 1x507x4
-    nms.confidenceOutputFeatureName = "confidence"
-    nms.coordinatesOutputFeatureName = "coordinates"
-    nms.iouThresholdInputFeatureName = "iouThreshold"
-    nms.confidenceThresholdInputFeatureName = "confidenceThreshold"
-    nms.iouThreshold = 0.45
-    nms.confidenceThreshold = 0.25
-    nms.pickTop.perClass = True
-    nms.stringClassLabels.vector.extend(names.values())
-    nms_model = ct.models.MLModel(nms_spec)
-
-    # 4. Pipeline models together
-    pipeline = ct.models.pipeline.Pipeline(
-        input_features=[
-            ("image", ct.models.datatypes.Array(3, ny, nx)),
-            ("iouThreshold", ct.models.datatypes.Double()),
-            ("confidenceThreshold", ct.models.datatypes.Double()),
-        ],
-        output_features=["confidence", "coordinates"],
-    )
-    pipeline.add_model(model)
-    pipeline.add_model(nms_model)
-
-    # Correct datatypes
-    pipeline.spec.description.input[0].ParseFromString(model._spec.description.input[0].SerializeToString())
-    pipeline.spec.description.output[0].ParseFromString(nms_model._spec.description.output[0].SerializeToString())
-    pipeline.spec.description.output[1].ParseFromString(nms_model._spec.description.output[1].SerializeToString())
-
-    # Update metadata
-    pipeline.spec.specificationVersion = 5
-    pipeline.spec.description.metadata.versionString = "https://github.com/ultralytics/yolov5"
-    pipeline.spec.description.metadata.shortDescription = "https://github.com/ultralytics/yolov5"
-    pipeline.spec.description.metadata.author = "glenn.jocher@ultralytics.com"
-    pipeline.spec.description.metadata.license = "https://github.com/ultralytics/yolov5/blob/master/LICENSE"
-    pipeline.spec.description.metadata.userDefined.update(
-        {
-            "classes": ",".join(names.values()),
-            "iou_threshold": str(nms.iouThreshold),
-            "confidence_threshold": str(nms.confidenceThreshold),
-        }
-    )
-
-    # Save the model
-    model = ct.models.MLModel(pipeline.spec, weights_dir=weights_dir)
-    model.input_description["image"] = "Input image"
-    model.input_description["iouThreshold"] = f"(optional) IOU Threshold override (default: {nms.iouThreshold})"
-    model.input_description["confidenceThreshold"] = (
-        f"(optional) Confidence Threshold override (default: {nms.confidenceThreshold})"
-    )
-    model.output_description["confidence"] = 'Boxes × Class confidence (see user-defined metadata "classes")'
-    model.output_description["coordinates"] = "Boxes × [x, y, width, height] (relative to image size)"
-    model.save(f)  # pipelined
-    print(f"{prefix} pipeline success ({time.time() - t:.2f}s), saved as {f} ({file_size(f):.1f} MB)")
-
-
-@smart_inference_mode()
-def run(
-    data=ROOT / "data/coco128.yaml",  # 'dataset.yaml path'
-    weights=ROOT / "yolov5s.pt",  # weights path
-    imgsz=(640, 640),  # image (height, width)
-    batch_size=1,  # batch size
-    device="cpu",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
-    include=("torchscript", "onnx"),  # include formats
-    half=False,  # FP16 half-precision export
-    inplace=False,  # set YOLOv5 Detect() inplace=True
-    keras=False,  # use Keras
-    optimize=False,  # TorchScript: optimize for mobile
-    int8=False,  # CoreML/TF INT8 quantization
-    per_tensor=False,  # TF per tensor quantization
-    dynamic=False,  # ONNX/TF/TensorRT: dynamic axes
-    simplify=False,  # ONNX: simplify model
-    mlmodel=False,  # CoreML: Export in *.mlmodel format
-    opset=12,  # ONNX: opset version
-    verbose=False,  # TensorRT: verbose log
-    workspace=4,  # TensorRT: workspace size (GB)
-    nms=False,  # TF: add NMS to model
-    agnostic_nms=False,  # TF: add agnostic NMS to model
-    topk_per_class=100,  # TF.js NMS: topk per class to keep
-    topk_all=100,  # TF.js NMS: topk for all classes to keep
-    iou_thres=0.45,  # TF.js NMS: IoU threshold
-    conf_thres=0.25,  # TF.js NMS: confidence threshold
-):
-    """
-    Exports a YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow.
-
-    Args:
-        data (str | Path): Path to the dataset YAML configuration file. Default is 'data/coco128.yaml'.
-        weights (str | Path): Path to the pretrained model weights file. Default is 'yolov5s.pt'.
-        imgsz (tuple): Image size as (height, width). Default is (640, 640).
-        batch_size (int): Batch size for exporting the model. Default is 1.
-        device (str): Device to run the export on, e.g., '0' for GPU, 'cpu' for CPU. Default is 'cpu'.
-        include (tuple): Formats to include in the export. Default is ('torchscript', 'onnx').
-        half (bool): Flag to export model with FP16 half-precision. Default is False.
-        inplace (bool): Set the YOLOv5 Detect() module inplace=True. Default is False.
-        keras (bool): Flag to use Keras for TensorFlow SavedModel export. Default is False.
-        optimize (bool): Optimize TorchScript model for mobile deployment. Default is False.
-        int8 (bool): Apply INT8 quantization for CoreML or TensorFlow models. Default is False.
-        per_tensor (bool): Apply per tensor quantization for TensorFlow models. Default is False.
-        dynamic (bool): Enable dynamic axes for ONNX, TensorFlow, or TensorRT exports. Default is False.
-        simplify (bool): Simplify the ONNX model during export. Default is False.
-        opset (int): ONNX opset version. Default is 12.
-        verbose (bool): Enable verbose logging for TensorRT export. Default is False.
-        workspace (int): TensorRT workspace size in GB. Default is 4.
-        nms (bool): Add non-maximum suppression (NMS) to the TensorFlow model. Default is False.
-        agnostic_nms (bool): Add class-agnostic NMS to the TensorFlow model. Default is False.
-        topk_per_class (int): Top-K boxes per class to keep for TensorFlow.js NMS. Default is 100.
-        topk_all (int): Top-K boxes for all classes to keep for TensorFlow.js NMS. Default is 100.
-        iou_thres (float): IoU threshold for NMS. Default is 0.45.
-        conf_thres (float): Confidence threshold for NMS. Default is 0.25.
-        mlmodel (bool): Flag to use *.mlmodel for CoreML export. Default is False.
-
-    Returns:
-        None
-
-    Notes:
-        - Model export is based on the specified formats in the 'include' argument.
-        - Be cautious of combinations where certain flags are mutually exclusive, such as `--half` and `--dynamic`.
-
-    Example:
-        ```python
-        run(
-            data="data/coco128.yaml",
-            weights="yolov5s.pt",
-            imgsz=(640, 640),
-            batch_size=1,
-            device="cpu",
-            include=("torchscript", "onnx"),
-            half=False,
-            inplace=False,
-            keras=False,
-            optimize=False,
-            int8=False,
-            per_tensor=False,
-            dynamic=False,
-            simplify=False,
-            opset=12,
-            verbose=False,
-            mlmodel=False,
-            workspace=4,
-            nms=False,
-            agnostic_nms=False,
-            topk_per_class=100,
-            topk_all=100,
-            iou_thres=0.45,
-            conf_thres=0.25,
-        )
-        ```
-    """
-    t = time.time()
-    include = [x.lower() for x in include]  # to lowercase
-    fmts = tuple(export_formats()["Argument"][1:])  # --include arguments
-    flags = [x in include for x in fmts]
-    assert sum(flags) == len(include), f"ERROR: Invalid --include {include}, valid --include arguments are {fmts}"
-    jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags  # export booleans
-    file = Path(url2file(weights) if str(weights).startswith(("http:/", "https:/")) else weights)  # PyTorch weights
-
-    # Load PyTorch model
-    device = select_device(device)
-    if half:
-        assert device.type != "cpu" or coreml, "--half only compatible with GPU export, i.e. use --device 0"
-        assert not dynamic, "--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both"
-    model = attempt_load(weights, device=device, inplace=True, fuse=True)  # load FP32 model
-
-    # Checks
-    imgsz *= 2 if len(imgsz) == 1 else 1  # expand
-    if optimize:
-        assert device.type == "cpu", "--optimize not compatible with cuda devices, i.e. use --device cpu"
-
-    # Input
-    gs = int(max(model.stride))  # grid size (max stride)
-    imgsz = [check_img_size(x, gs) for x in imgsz]  # verify img_size are gs-multiples
-    im = torch.zeros(batch_size, 3, *imgsz).to(device)  # image size(1,3,320,192) BCHW iDetection
-
-    # Update model
-    model.eval()
-    for k, m in model.named_modules():
-        if isinstance(m, Detect):
-            m.inplace = inplace
-            m.dynamic = dynamic
-            m.export = True
-
-    for _ in range(2):
-        y = model(im)  # dry runs
-    if half and not coreml:
-        im, model = im.half(), model.half()  # to FP16
-    shape = tuple((y[0] if isinstance(y, tuple) else y).shape)  # model output shape
-    metadata = {"stride": int(max(model.stride)), "names": model.names}  # model metadata
-    LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")
-
-    # Exports
-    f = [""] * len(fmts)  # exported filenames
-    warnings.filterwarnings(action="ignore", category=torch.jit.TracerWarning)  # suppress TracerWarning
-    if jit:  # TorchScript
-        f[0], _ = export_torchscript(model, im, file, optimize)
-    if engine:  # TensorRT required before ONNX
-        f[1], _ = export_engine(model, im, file, half, dynamic, simplify, workspace, verbose)
-    if onnx or xml:  # OpenVINO requires ONNX
-        f[2], _ = export_onnx(model, im, file, opset, dynamic, simplify)
-    if xml:  # OpenVINO
-        f[3], _ = export_openvino(file, metadata, half, int8, data)
-    if coreml:  # CoreML
-        f[4], ct_model = export_coreml(model, im, file, int8, half, nms, mlmodel)
-        if nms:
-            pipeline_coreml(ct_model, im, file, model.names, y, mlmodel)
-    if any((saved_model, pb, tflite, edgetpu, tfjs)):  # TensorFlow formats
-        assert not tflite or not tfjs, "TFLite and TF.js models must be exported separately, please pass only one type."
-        assert not isinstance(model, ClassificationModel), "ClassificationModel export to TF formats not yet supported."
-        f[5], s_model = export_saved_model(
-            model.cpu(),
-            im,
-            file,
-            dynamic,
-            tf_nms=nms or agnostic_nms or tfjs,
-            agnostic_nms=agnostic_nms or tfjs,
-            topk_per_class=topk_per_class,
-            topk_all=topk_all,
-            iou_thres=iou_thres,
-            conf_thres=conf_thres,
-            keras=keras,
-        )
-        if pb or tfjs:  # pb prerequisite to tfjs
-            f[6], _ = export_pb(s_model, file)
-        if tflite or edgetpu:
-            f[7], _ = export_tflite(
-                s_model, im, file, int8 or edgetpu, per_tensor, data=data, nms=nms, agnostic_nms=agnostic_nms
-            )
-            if edgetpu:
-                f[8], _ = export_edgetpu(file)
-            add_tflite_metadata(f[8] or f[7], metadata, num_outputs=len(s_model.outputs))
-        if tfjs:
-            f[9], _ = export_tfjs(file, int8)
-    if paddle:  # PaddlePaddle
-        f[10], _ = export_paddle(model, im, file, metadata)
-
-    # Finish
-    f = [str(x) for x in f if x]  # filter out '' and None
-    if any(f):
-        cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel))  # type
-        det &= not seg  # segmentation models inherit from SegmentationModel(DetectionModel)
-        dir = Path("segment" if seg else "classify" if cls else "")
-        h = "--half" if half else ""  # --half FP16 inference arg
-        s = (
-            "# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference"
-            if cls
-            else "# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference"
-            if seg
-            else ""
-        )
-        LOGGER.info(
-            f'\nExport complete ({time.time() - t:.1f}s)'
-            f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
-            f"\nDetect:          python {dir / ('detect.py' if det else 'predict.py')} --weights {f[-1]} {h}"
-            f"\nValidate:        python {dir / 'val.py'} --weights {f[-1]} {h}"
-            f"\nPyTorch Hub:     model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')  {s}"
-            f'\nVisualize:       https://netron.app'
-        )
-    return f  # return list of exported files/dirs
-
-
-def parse_opt(known=False):
-    """
-    Parse command-line options for YOLOv5 model export configurations.
-
-    Args:
-        known (bool): If True, uses `argparse.ArgumentParser.parse_known_args`; otherwise, uses `argparse.ArgumentParser.parse_args`.
-                      Default is False.
-
-    Returns:
-        argparse.Namespace: Object containing parsed command-line arguments.
-
-    Example:
-        ```python
-        opts = parse_opt()
-        print(opts.data)
-        print(opts.weights)
-        ```
-    """
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
-    parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model.pt path(s)")
-    parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640, 640], help="image (h, w)")
-    parser.add_argument("--batch-size", type=int, default=1, help="batch size")
-    parser.add_argument("--device", default="cpu", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
-    parser.add_argument("--half", action="store_true", help="FP16 half-precision export")
-    parser.add_argument("--inplace", action="store_true", help="set YOLOv5 Detect() inplace=True")
-    parser.add_argument("--keras", action="store_true", help="TF: use Keras")
-    parser.add_argument("--optimize", action="store_true", help="TorchScript: optimize for mobile")
-    parser.add_argument("--int8", action="store_true", help="CoreML/TF/OpenVINO INT8 quantization")
-    parser.add_argument("--per-tensor", action="store_true", help="TF per-tensor quantization")
-    parser.add_argument("--dynamic", action="store_true", help="ONNX/TF/TensorRT: dynamic axes")
-    parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model")
-    parser.add_argument("--mlmodel", action="store_true", help="CoreML: Export in *.mlmodel format")
-    parser.add_argument("--opset", type=int, default=17, help="ONNX: opset version")
-    parser.add_argument("--verbose", action="store_true", help="TensorRT: verbose log")
-    parser.add_argument("--workspace", type=int, default=4, help="TensorRT: workspace size (GB)")
-    parser.add_argument("--nms", action="store_true", help="TF: add NMS to model")
-    parser.add_argument("--agnostic-nms", action="store_true", help="TF: add agnostic NMS to model")
-    parser.add_argument("--topk-per-class", type=int, default=100, help="TF.js NMS: topk per class to keep")
-    parser.add_argument("--topk-all", type=int, default=100, help="TF.js NMS: topk for all classes to keep")
-    parser.add_argument("--iou-thres", type=float, default=0.45, help="TF.js NMS: IoU threshold")
-    parser.add_argument("--conf-thres", type=float, default=0.25, help="TF.js NMS: confidence threshold")
-    parser.add_argument(
-        "--include",
-        nargs="+",
-        default=["torchscript"],
-        help="torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle",
-    )
-    opt = parser.parse_known_args()[0] if known else parser.parse_args()
-    print_args(vars(opt))
-    return opt
-
-
-def main(opt):
-    """Run(**vars(opt))  # Execute the run function with parsed options."""
-    for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
-        run(**vars(opt))
-
-
-if __name__ == "__main__":
-    opt = parse_opt()
-    main(opt)