Reference for ultralytics/engine/exporter.py

def __call__(self, model=None) -> str:
    """Return list of exported files/dirs after running callbacks."""
    t = time.time()
    fmt = self.args.format.lower()  # to lowercase
    if fmt in {"tensorrt", "trt"}:  # 'engine' aliases
        fmt = "engine"
    if fmt in {"mlmodel", "mlpackage", "mlprogram", "apple", "ios", "coreml"}:  # 'coreml' aliases
        fmt = "coreml"
    fmts_dict = export_formats()
    fmts = tuple(fmts_dict["Argument"][1:])  # available export formats
    if fmt not in fmts:
        import difflib

        # Get the closest match if format is invalid
        matches = difflib.get_close_matches(fmt, fmts, n=1, cutoff=0.6)  # 60% similarity required to match
        if not matches:
            raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
        LOGGER.warning(f"Invalid export format='{fmt}', updating to format='{matches[0]}'")
        fmt = matches[0]
    flags = [x == fmt for x in fmts]
    if sum(flags) != 1:
        raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
    (jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, mnn, ncnn, imx, rknn) = (
        flags  # export booleans
    )

    is_tf_format = any((saved_model, pb, tflite, edgetpu, tfjs))

    # Device
    dla = None
    if engine and self.args.device is None:
        LOGGER.warning("TensorRT requires GPU export, automatically assigning device=0")
        self.args.device = "0"
    if engine and "dla" in str(self.args.device):  # convert int/list to str first
        dla = self.args.device.rsplit(":", 1)[-1]
        self.args.device = "0"  # update device to "0"
        assert dla in {"0", "1"}, f"Expected self.args.device='dla:0' or 'dla:1, but got {self.args.device}."
    if imx and self.args.device is None and torch.cuda.is_available():
        LOGGER.warning("Exporting on CPU while CUDA is available, setting device=0 for faster export on GPU.")
        self.args.device = "0"  # update device to "0"
    self.device = select_device("cpu" if self.args.device is None else self.args.device)

    # Argument compatibility checks
    fmt_keys = fmts_dict["Arguments"][flags.index(True) + 1]
    validate_args(fmt, self.args, fmt_keys)
    if imx:
        if not self.args.int8:
            LOGGER.warning("IMX export requires int8=True, setting int8=True.")
            self.args.int8 = True
        if not self.args.nms:
            LOGGER.warning("IMX export requires nms=True, setting nms=True.")
            self.args.nms = True
        if model.task not in {"detect", "pose"}:
            raise ValueError("IMX export only supported for detection and pose estimation models.")
    if not hasattr(model, "names"):
        model.names = default_class_names()
    model.names = check_class_names(model.names)
    if self.args.half and self.args.int8:
        LOGGER.warning("half=True and int8=True are mutually exclusive, setting half=False.")
        self.args.half = False
    if self.args.half and onnx and self.device.type == "cpu":
        LOGGER.warning("half=True only compatible with GPU export, i.e. use device=0")
        self.args.half = False
    self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2)  # check image size
    if self.args.optimize:
        assert not ncnn, "optimize=True not compatible with format='ncnn', i.e. use optimize=False"
        assert self.device.type == "cpu", "optimize=True not compatible with cuda devices, i.e. use device='cpu'"
    if rknn:
        if not self.args.name:
            LOGGER.warning(
                "Rockchip RKNN export requires a missing 'name' arg for processor type. "
                "Using default name='rk3588'."
            )
            self.args.name = "rk3588"
        self.args.name = self.args.name.lower()
        assert self.args.name in RKNN_CHIPS, (
            f"Invalid processor name '{self.args.name}' for Rockchip RKNN export. Valid names are {RKNN_CHIPS}."
        )
    if self.args.int8 and tflite:
        assert not getattr(model, "end2end", False), "TFLite INT8 export not supported for end2end models."
    if self.args.nms:
        assert not isinstance(model, ClassificationModel), "'nms=True' is not valid for classification models."
        assert not tflite or not ARM64 or not LINUX, "TFLite export with NMS unsupported on ARM64 Linux"
        if getattr(model, "end2end", False):
            LOGGER.warning("'nms=True' is not available for end2end models. Forcing 'nms=False'.")
            self.args.nms = False
        self.args.conf = self.args.conf or 0.25  # set conf default value for nms export
    if (engine or self.args.nms) and self.args.dynamic and self.args.batch == 1:
        LOGGER.warning(
            f"'dynamic=True' model with '{'nms=True' if self.args.nms else 'format=engine'}' requires max batch size, i.e. 'batch=16'"
        )
    if edgetpu:
        if not LINUX or ARM64:
            raise SystemError(
                "Edge TPU export only supported on non-aarch64 Linux. See https://coral.ai/docs/edgetpu/compiler"
            )
        elif self.args.batch != 1:  # see github.com/ultralytics/ultralytics/pull/13420
            LOGGER.warning("Edge TPU export requires batch size 1, setting batch=1.")
            self.args.batch = 1
    if isinstance(model, WorldModel):
        LOGGER.warning(
            "YOLOWorld (original version) export is not supported to any format. "
            "YOLOWorldv2 models (i.e. 'yolov8s-worldv2.pt') only support export to "
            "(torchscript, onnx, openvino, engine, coreml) formats. "
            "See https://docs.ultralytics.com/models/yolo-world for details."
        )
        model.clip_model = None  # openvino int8 export error: https://github.com/ultralytics/ultralytics/pull/18445
    if self.args.int8 and not self.args.data:
        self.args.data = DEFAULT_CFG.data or TASK2DATA[getattr(model, "task", "detect")]  # assign default data
        LOGGER.warning(
            f"INT8 export requires a missing 'data' arg for calibration. Using default 'data={self.args.data}'."
        )
    if tfjs and (ARM64 and LINUX):
        raise SystemError("TF.js exports are not currently supported on ARM64 Linux")
    # Recommend OpenVINO if export and Intel CPU
    if SETTINGS.get("openvino_msg"):
        if is_intel():
            LOGGER.info(
                "💡 ProTip: Export to OpenVINO format for best performance on Intel hardware."
                " Learn more at https://docs.ultralytics.com/integrations/openvino/"
            )
        SETTINGS["openvino_msg"] = False

    # Input
    im = torch.zeros(self.args.batch, model.yaml.get("channels", 3), *self.imgsz).to(self.device)
    file = Path(
        getattr(model, "pt_path", None) or getattr(model, "yaml_file", None) or model.yaml.get("yaml_file", "")
    )
    if file.suffix in {".yaml", ".yml"}:
        file = Path(file.name)

    # Update model
    model = deepcopy(model).to(self.device)
    for p in model.parameters():
        p.requires_grad = False
    model.eval()
    model.float()
    model = model.fuse()

    if imx:
        from ultralytics.utils.torch_utils import FXModel

        model = FXModel(model)
    for m in model.modules():
        if isinstance(m, Classify):
            m.export = True
        if isinstance(m, (Detect, RTDETRDecoder)):  # includes all Detect subclasses like Segment, Pose, OBB
            m.dynamic = self.args.dynamic
            m.export = True
            m.format = self.args.format
            m.max_det = self.args.max_det
            m.xyxy = self.args.nms and not coreml
            if hasattr(model, "pe") and hasattr(m, "fuse"):  # for YOLOE models
                m.fuse(model.pe.to(self.device))
        elif isinstance(m, C2f) and not is_tf_format:
            # EdgeTPU does not support FlexSplitV while split provides cleaner ONNX graph
            m.forward = m.forward_split
        if isinstance(m, Detect) and imx:
            from ultralytics.utils.tal import make_anchors

            m.anchors, m.strides = (
                x.transpose(0, 1)
                for x in make_anchors(
                    torch.cat([s / m.stride.unsqueeze(-1) for s in self.imgsz], dim=1), m.stride, 0.5
                )
            )

    y = None
    for _ in range(2):  # dry runs
        y = NMSModel(model, self.args)(im) if self.args.nms and not coreml and not imx else model(im)
    if self.args.half and onnx and self.device.type != "cpu":
        im, model = im.half(), model.half()  # to FP16

    # Filter warnings
    warnings.filterwarnings("ignore", category=torch.jit.TracerWarning)  # suppress TracerWarning
    warnings.filterwarnings("ignore", category=UserWarning)  # suppress shape prim::Constant missing ONNX warning
    warnings.filterwarnings("ignore", category=DeprecationWarning)  # suppress CoreML np.bool deprecation warning

    # Assign
    self.im = im
    self.model = model
    self.file = file
    self.output_shape = (
        tuple(y.shape)
        if isinstance(y, torch.Tensor)
        else tuple(tuple(x.shape if isinstance(x, torch.Tensor) else []) for x in y)
    )
    self.pretty_name = Path(self.model.yaml.get("yaml_file", self.file)).stem.replace("yolo", "YOLO")
    data = model.args["data"] if hasattr(model, "args") and isinstance(model.args, dict) else ""
    description = f"Ultralytics {self.pretty_name} model {f'trained on {data}' if data else ''}"
    self.metadata = {
        "description": description,
        "author": "Ultralytics",
        "date": datetime.now().isoformat(),
        "version": __version__,
        "license": "AGPL-3.0 License (https://ultralytics.com/license)",
        "docs": "https://docs.ultralytics.com",
        "stride": int(max(model.stride)),
        "task": model.task,
        "batch": self.args.batch,
        "imgsz": self.imgsz,
        "names": model.names,
        "args": {k: v for k, v in self.args if k in fmt_keys},
        "channels": model.yaml.get("channels", 3),
    }  # model metadata
    if dla is not None:
        self.metadata["dla"] = dla  # make sure `AutoBackend` uses correct dla device if it has one
    if model.task == "pose":
        self.metadata["kpt_shape"] = model.model[-1].kpt_shape

    LOGGER.info(
        f"\n{colorstr('PyTorch:')} starting from '{file}' with input shape {tuple(im.shape)} BCHW and "
        f"output shape(s) {self.output_shape} ({file_size(file):.1f} MB)"
    )
    self.run_callbacks("on_export_start")
    # Exports
    f = [""] * len(fmts)  # exported filenames
    if jit or ncnn:  # TorchScript
        f[0], _ = self.export_torchscript()
    if engine:  # TensorRT required before ONNX
        f[1], _ = self.export_engine(dla=dla)
    if onnx:  # ONNX
        f[2], _ = self.export_onnx()
    if xml:  # OpenVINO
        f[3], _ = self.export_openvino()
    if coreml:  # CoreML
        f[4], _ = self.export_coreml()
    if is_tf_format:  # TensorFlow formats
        self.args.int8 |= edgetpu
        f[5], keras_model = self.export_saved_model()
        if pb or tfjs:  # pb prerequisite to tfjs
            f[6], _ = self.export_pb(keras_model=keras_model)
        if tflite:
            f[7], _ = self.export_tflite()
        if edgetpu:
            f[8], _ = self.export_edgetpu(tflite_model=Path(f[5]) / f"{self.file.stem}_full_integer_quant.tflite")
        if tfjs:
            f[9], _ = self.export_tfjs()
    if paddle:  # PaddlePaddle
        f[10], _ = self.export_paddle()
    if mnn:  # MNN
        f[11], _ = self.export_mnn()
    if ncnn:  # NCNN
        f[12], _ = self.export_ncnn()
    if imx:
        f[13], _ = self.export_imx()
    if rknn:
        f[14], _ = self.export_rknn()

    # Finish
    f = [str(x) for x in f if x]  # filter out '' and None
    if any(f):
        f = str(Path(f[-1]))
        square = self.imgsz[0] == self.imgsz[1]
        s = (
            ""
            if square
            else f"WARNING ⚠️ non-PyTorch val requires square images, 'imgsz={self.imgsz}' will not "
            f"work. Use export 'imgsz={max(self.imgsz)}' if val is required."
        )
        imgsz = self.imgsz[0] if square else str(self.imgsz)[1:-1].replace(" ", "")
        predict_data = f"data={data}" if model.task == "segment" and pb else ""
        q = "int8" if self.args.int8 else "half" if self.args.half else ""  # quantization
        LOGGER.info(
            f"\nExport complete ({time.time() - t:.1f}s)"
            f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
            f"\nPredict:         yolo predict task={model.task} model={f} imgsz={imgsz} {q} {predict_data}"
            f"\nValidate:        yolo val task={model.task} model={f} imgsz={imgsz} data={data} {q} {s}"
            f"\nVisualize:       https://netron.app"
        )

    self.run_callbacks("on_export_end")
    return f  # return list of exported files/dirs

def add_callback(self, event: str, callback):
    """Append the given callback to the specified event."""
    self.callbacks[event].append(callback)

@try_export
def export_coreml(self, prefix=colorstr("CoreML:")):
    """Export YOLO model to CoreML format."""
    mlmodel = self.args.format.lower() == "mlmodel"  # legacy *.mlmodel export format requested
    check_requirements("coremltools>=8.0")
    import coremltools as ct  # noqa

    LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
    assert not WINDOWS, "CoreML export is not supported on Windows, please run on macOS or Linux."
    assert self.args.batch == 1, "CoreML batch sizes > 1 are not supported. Please retry at 'batch=1'."
    f = self.file.with_suffix(".mlmodel" if mlmodel else ".mlpackage")
    if f.is_dir():
        shutil.rmtree(f)

    bias = [0.0, 0.0, 0.0]
    scale = 1 / 255
    classifier_config = None
    if self.model.task == "classify":
        classifier_config = ct.ClassifierConfig(list(self.model.names.values()))
        model = self.model
    elif self.model.task == "detect":
        model = IOSDetectModel(self.model, self.im) if self.args.nms else self.model
    else:
        if self.args.nms:
            LOGGER.warning(f"{prefix} 'nms=True' is only available for Detect models like 'yolo11n.pt'.")
            # TODO CoreML Segment and Pose model pipelining
        model = self.model
    ts = torch.jit.trace(model.eval(), self.im, strict=False)  # TorchScript model

    # Based on apple's documentation it is better to leave out the minimum_deployment target and let that get set
    # Internally based on the model conversion and output type.
    # Setting minimum_depoloyment_target >= iOS16 will require setting compute_precision=ct.precision.FLOAT32.
    # iOS16 adds in better support for FP16, but none of the CoreML NMS specifications handle FP16 as input.
    ct_model = ct.convert(
        ts,
        inputs=[ct.ImageType("image", shape=self.im.shape, scale=scale, bias=bias)],  # expects ct.TensorType
        classifier_config=classifier_config,
        convert_to="neuralnetwork" if mlmodel else "mlprogram",
    )
    bits, mode = (8, "kmeans") if self.args.int8 else (16, "linear") if self.args.half else (32, None)
    if bits < 32:
        if "kmeans" in mode:
            check_requirements("scikit-learn")  # scikit-learn package required for k-means quantization
        if mlmodel:
            ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
        elif bits == 8:  # mlprogram already quantized to FP16
            import coremltools.optimize.coreml as cto

            op_config = cto.OpPalettizerConfig(mode="kmeans", nbits=bits, weight_threshold=512)
            config = cto.OptimizationConfig(global_config=op_config)
            ct_model = cto.palettize_weights(ct_model, config=config)
    if self.args.nms and self.model.task == "detect":
        if mlmodel:
            weights_dir = None
        else:
            ct_model.save(str(f))  # save otherwise weights_dir does not exist
            weights_dir = str(f / "Data/com.apple.CoreML/weights")
        ct_model = self._pipeline_coreml(ct_model, weights_dir=weights_dir)

    m = self.metadata  # metadata dict
    ct_model.short_description = m.pop("description")
    ct_model.author = m.pop("author")
    ct_model.license = m.pop("license")
    ct_model.version = m.pop("version")
    ct_model.user_defined_metadata.update({k: str(v) for k, v in m.items()})
    if self.model.task == "classify":
        ct_model.user_defined_metadata.update({"com.apple.coreml.model.preview.type": "imageClassifier"})

    try:
        ct_model.save(str(f))  # save *.mlpackage
    except Exception as e:
        LOGGER.warning(
            f"{prefix} CoreML export to *.mlpackage failed ({e}), reverting to *.mlmodel export. "
            f"Known coremltools Python 3.11 and Windows bugs https://github.com/apple/coremltools/issues/1928."
        )
        f = f.with_suffix(".mlmodel")
        ct_model.save(str(f))
    return f, ct_model

@try_export
def export_edgetpu(self, tflite_model="", prefix=colorstr("Edge TPU:")):
    """Export YOLO model to Edge TPU format https://coral.ai/docs/edgetpu/models-intro/."""
    cmd = "edgetpu_compiler --version"
    help_url = "https://coral.ai/docs/edgetpu/compiler/"
    assert LINUX, f"export only supported on Linux. See {help_url}"
    if subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=True).returncode != 0:
        LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}")
        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 is_sudo_available() else c.replace("sudo ", ""), shell=True, check=True)
    ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().rsplit(maxsplit=1)[-1]

    LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...")
    f = str(tflite_model).replace(".tflite", "_edgetpu.tflite")  # Edge TPU model

    cmd = (
        "edgetpu_compiler "
        f'--out_dir "{Path(f).parent}" '
        "--show_operations "
        "--search_delegate "
        "--delegate_search_step 30 "
        "--timeout_sec 180 "
        f'"{tflite_model}"'
    )
    LOGGER.info(f"{prefix} running '{cmd}'")
    subprocess.run(cmd, shell=True)
    self._add_tflite_metadata(f)
    return f, None

@try_export
def export_engine(self, dla=None, prefix=colorstr("TensorRT:")):
    """Export YOLO model to TensorRT format https://developer.nvidia.com/tensorrt."""
    assert self.im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. use 'device=0'"
    f_onnx, _ = self.export_onnx()  # run before TRT import https://github.com/ultralytics/ultralytics/issues/7016

    try:
        import tensorrt as trt  # noqa
    except ImportError:
        if LINUX:
            check_requirements("tensorrt>7.0.0,!=10.1.0")
        import tensorrt as trt  # noqa
    check_version(trt.__version__, ">=7.0.0", hard=True)
    check_version(trt.__version__, "!=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")

    # Setup and checks
    LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")
    assert Path(f_onnx).exists(), f"failed to export ONNX file: {f_onnx}"
    f = self.file.with_suffix(".engine")  # TensorRT engine file
    export_engine(
        f_onnx,
        f,
        self.args.workspace,
        self.args.half,
        self.args.int8,
        self.args.dynamic,
        self.im.shape,
        dla=dla,
        dataset=self.get_int8_calibration_dataloader(prefix) if self.args.int8 else None,
        metadata=self.metadata,
        verbose=self.args.verbose,
        prefix=prefix,
    )

    return f, None

@try_export
def export_imx(self, prefix=colorstr("IMX:")):
    """Export YOLO model to IMX format."""
    gptq = False
    assert LINUX, (
        "export only supported on Linux. "
        "See https://developer.aitrios.sony-semicon.com/en/raspberrypi-ai-camera/documentation/imx500-converter"
    )
    if getattr(self.model, "end2end", False):
        raise ValueError("IMX export is not supported for end2end models.")
    check_requirements(
        ("model-compression-toolkit>=2.4.1", "sony-custom-layers>=0.3.0", "edge-mdt-tpc>=1.1.0", "pydantic<=2.11.7")
    )
    check_requirements("imx500-converter[pt]>=3.16.1")  # Separate requirements for imx500-converter
    check_requirements("mct-quantizers>=1.6.0")  # Separate for compatibility with model-compression-toolkit

    import model_compression_toolkit as mct
    import onnx
    from edgemdt_tpc import get_target_platform_capabilities
    from sony_custom_layers.pytorch import multiclass_nms_with_indices

    LOGGER.info(f"\n{prefix} starting export with model_compression_toolkit {mct.__version__}...")

    # Install Java>=17
    try:
        java_output = subprocess.run(["java", "--version"], check=True, capture_output=True).stdout.decode()
        version_match = re.search(r"(?:openjdk|java) (\d+)", java_output)
        java_version = int(version_match.group(1)) if version_match else 0
        assert java_version >= 17, "Java version too old"
    except (FileNotFoundError, subprocess.CalledProcessError, AssertionError):
        cmd = (["sudo"] if is_sudo_available() else []) + ["apt", "install", "-y", "openjdk-21-jre"]
        subprocess.run(cmd, check=True)

    def representative_dataset_gen(dataloader=self.get_int8_calibration_dataloader(prefix)):
        for batch in dataloader:
            img = batch["img"]
            img = img / 255.0
            yield [img]

    tpc = get_target_platform_capabilities(tpc_version="4.0", device_type="imx500")

    bit_cfg = mct.core.BitWidthConfig()
    if "C2PSA" in self.model.__str__():  # YOLO11
        if self.model.task == "detect":
            layer_names = ["sub", "mul_2", "add_14", "cat_21"]
            weights_memory = 2585350.2439
            n_layers = 238  # 238 layers for fused YOLO11n
        elif self.model.task == "pose":
            layer_names = ["sub", "mul_2", "add_14", "cat_22", "cat_23", "mul_4", "add_15"]
            weights_memory = 2437771.67
            n_layers = 257  # 257 layers for fused YOLO11n-pose
    else:  # YOLOv8
        if self.model.task == "detect":
            layer_names = ["sub", "mul", "add_6", "cat_17"]
            weights_memory = 2550540.8
            n_layers = 168  # 168 layers for fused YOLOv8n
        elif self.model.task == "pose":
            layer_names = ["add_7", "mul_2", "cat_19", "mul", "sub", "add_6", "cat_18"]
            weights_memory = 2482451.85
            n_layers = 187  # 187 layers for fused YOLO11n-pose

    # Check if the model has the expected number of layers
    if len(list(self.model.modules())) != n_layers:
        raise ValueError("IMX export only supported for YOLOv8n and YOLO11n models.")

    for layer_name in layer_names:
        bit_cfg.set_manual_activation_bit_width([mct.core.common.network_editors.NodeNameFilter(layer_name)], 16)

    config = mct.core.CoreConfig(
        mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=10),
        quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True),
        bit_width_config=bit_cfg,
    )

    resource_utilization = mct.core.ResourceUtilization(weights_memory=weights_memory)

    quant_model = (
        mct.gptq.pytorch_gradient_post_training_quantization(  # Perform Gradient-Based Post Training Quantization
            model=self.model,
            representative_data_gen=representative_dataset_gen,
            target_resource_utilization=resource_utilization,
            gptq_config=mct.gptq.get_pytorch_gptq_config(
                n_epochs=1000, use_hessian_based_weights=False, use_hessian_sample_attention=False
            ),
            core_config=config,
            target_platform_capabilities=tpc,
        )[0]
        if gptq
        else mct.ptq.pytorch_post_training_quantization(  # Perform post training quantization
            in_module=self.model,
            representative_data_gen=representative_dataset_gen,
            target_resource_utilization=resource_utilization,
            core_config=config,
            target_platform_capabilities=tpc,
        )[0]
    )

    class NMSWrapper(torch.nn.Module):
        """Wrap PyTorch Module with multiclass_nms layer from sony_custom_layers."""

        def __init__(
            self,
            model: torch.nn.Module,
            score_threshold: float = 0.001,
            iou_threshold: float = 0.7,
            max_detections: int = 300,
            task: str = "detect",
        ):
            """
            Initialize NMSWrapper with PyTorch Module and NMS parameters.

            Args:
                model (torch.nn.Module): Model instance.
                score_threshold (float): Score threshold for non-maximum suppression.
                iou_threshold (float): Intersection over union threshold for non-maximum suppression.
                max_detections (int): The number of detections to return.
                task (str): Task type, either 'detect' or 'pose'.
            """
            super().__init__()
            self.model = model
            self.score_threshold = score_threshold
            self.iou_threshold = iou_threshold
            self.max_detections = max_detections
            self.task = task

        def forward(self, images):
            """Forward pass with model inference and NMS post-processing."""
            # model inference
            outputs = self.model(images)

            boxes, scores = outputs[0], outputs[1]
            nms_outputs = multiclass_nms_with_indices(
                boxes=boxes,
                scores=scores,
                score_threshold=self.score_threshold,
                iou_threshold=self.iou_threshold,
                max_detections=self.max_detections,
            )
            if self.task == "pose":
                kpts = outputs[2]  # (bs, max_detections, kpts 17*3)
                out_kpts = torch.gather(kpts, 1, nms_outputs.indices.unsqueeze(-1).expand(-1, -1, kpts.size(-1)))
                return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, out_kpts
            return nms_outputs

    quant_model = NMSWrapper(
        model=quant_model,
        score_threshold=self.args.conf or 0.001,
        iou_threshold=self.args.iou,
        max_detections=self.args.max_det,
        task=self.model.task,
    ).to(self.device)

    f = Path(str(self.file).replace(self.file.suffix, "_imx_model"))
    f.mkdir(exist_ok=True)
    onnx_model = f / Path(str(self.file.name).replace(self.file.suffix, "_imx.onnx"))  # js dir
    mct.exporter.pytorch_export_model(
        model=quant_model, save_model_path=onnx_model, repr_dataset=representative_dataset_gen
    )

    model_onnx = onnx.load(onnx_model)  # load onnx model
    for k, v in self.metadata.items():
        meta = model_onnx.metadata_props.add()
        meta.key, meta.value = k, str(v)

    onnx.save(model_onnx, onnx_model)

    subprocess.run(
        ["imxconv-pt", "-i", str(onnx_model), "-o", str(f), "--no-input-persistency", "--overwrite-output"],
        check=True,
    )

    # Needed for imx models.
    with open(f / "labels.txt", "w", encoding="utf-8") as file:
        file.writelines([f"{name}\n" for _, name in self.model.names.items()])

    return f, None

@try_export
def export_mnn(self, prefix=colorstr("MNN:")):
    """Export YOLO model to MNN format using MNN https://github.com/alibaba/MNN."""
    f_onnx, _ = self.export_onnx()  # get onnx model first

    check_requirements("MNN>=2.9.6")
    import MNN  # noqa
    from MNN.tools import mnnconvert

    # Setup and checks
    LOGGER.info(f"\n{prefix} starting export with MNN {MNN.version()}...")
    assert Path(f_onnx).exists(), f"failed to export ONNX file: {f_onnx}"
    f = str(self.file.with_suffix(".mnn"))  # MNN model file
    args = ["", "-f", "ONNX", "--modelFile", f_onnx, "--MNNModel", f, "--bizCode", json.dumps(self.metadata)]
    if self.args.int8:
        args.extend(("--weightQuantBits", "8"))
    if self.args.half:
        args.append("--fp16")
    mnnconvert.convert(args)
    # remove scratch file for model convert optimize
    convert_scratch = Path(self.file.parent / ".__convert_external_data.bin")
    if convert_scratch.exists():
        convert_scratch.unlink()
    return f, None

@try_export
def export_ncnn(self, prefix=colorstr("NCNN:")):
    """Export YOLO model to NCNN format using PNNX https://github.com/pnnx/pnnx."""
    check_requirements("ncnn", cmds="--no-deps")  # no deps to avoid installing opencv-python
    import ncnn  # noqa

    LOGGER.info(f"\n{prefix} starting export with NCNN {ncnn.__version__}...")
    f = Path(str(self.file).replace(self.file.suffix, f"_ncnn_model{os.sep}"))
    f_ts = self.file.with_suffix(".torchscript")

    name = Path("pnnx.exe" if WINDOWS else "pnnx")  # PNNX filename
    pnnx = name if name.is_file() else (ROOT / name)
    if not pnnx.is_file():
        LOGGER.warning(
            f"{prefix} PNNX not found. Attempting to download binary file from "
            "https://github.com/pnnx/pnnx/.\nNote PNNX Binary file must be placed in current working directory "
            f"or in {ROOT}. See PNNX repo for full installation instructions."
        )
        system = "macos" if MACOS else "windows" if WINDOWS else "linux-aarch64" if ARM64 else "linux"
        try:
            release, assets = get_github_assets(repo="pnnx/pnnx")
            asset = [x for x in assets if f"{system}.zip" in x][0]
            assert isinstance(asset, str), "Unable to retrieve PNNX repo assets"  # i.e. pnnx-20240410-macos.zip
            LOGGER.info(f"{prefix} successfully found latest PNNX asset file {asset}")
        except Exception as e:
            release = "20240410"
            asset = f"pnnx-{release}-{system}.zip"
            LOGGER.warning(f"{prefix} PNNX GitHub assets not found: {e}, using default {asset}")
        unzip_dir = safe_download(f"https://github.com/pnnx/pnnx/releases/download/{release}/{asset}", delete=True)
        if check_is_path_safe(Path.cwd(), unzip_dir):  # avoid path traversal security vulnerability
            shutil.move(src=unzip_dir / name, dst=pnnx)  # move binary to ROOT
            pnnx.chmod(0o777)  # set read, write, and execute permissions for everyone
            shutil.rmtree(unzip_dir)  # delete unzip dir

    ncnn_args = [
        f"ncnnparam={f / 'model.ncnn.param'}",
        f"ncnnbin={f / 'model.ncnn.bin'}",
        f"ncnnpy={f / 'model_ncnn.py'}",
    ]

    pnnx_args = [
        f"pnnxparam={f / 'model.pnnx.param'}",
        f"pnnxbin={f / 'model.pnnx.bin'}",
        f"pnnxpy={f / 'model_pnnx.py'}",
        f"pnnxonnx={f / 'model.pnnx.onnx'}",
    ]

    cmd = [
        str(pnnx),
        str(f_ts),
        *ncnn_args,
        *pnnx_args,
        f"fp16={int(self.args.half)}",
        f"device={self.device.type}",
        f'inputshape="{[self.args.batch, 3, *self.imgsz]}"',
    ]
    f.mkdir(exist_ok=True)  # make ncnn_model directory
    LOGGER.info(f"{prefix} running '{' '.join(cmd)}'")
    subprocess.run(cmd, check=True)

    # Remove debug files
    pnnx_files = [x.rsplit("=", 1)[-1] for x in pnnx_args]
    for f_debug in ("debug.bin", "debug.param", "debug2.bin", "debug2.param", *pnnx_files):
        Path(f_debug).unlink(missing_ok=True)

    YAML.save(f / "metadata.yaml", self.metadata)  # add metadata.yaml
    return str(f), None

@try_export
def export_onnx(self, prefix=colorstr("ONNX:")):
    """Export YOLO model to ONNX format."""
    requirements = ["onnx>=1.12.0"]
    if self.args.simplify:
        requirements += ["onnxslim>=0.1.65", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
    check_requirements(requirements)
    import onnx  # noqa

    opset_version = self.args.opset or get_latest_opset()
    LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__} opset {opset_version}...")
    f = str(self.file.with_suffix(".onnx"))
    output_names = ["output0", "output1"] if isinstance(self.model, SegmentationModel) else ["output0"]
    dynamic = self.args.dynamic
    if dynamic:
        dynamic = {"images": {0: "batch", 2: "height", 3: "width"}}  # shape(1,3,640,640)
        if isinstance(self.model, SegmentationModel):
            dynamic["output0"] = {0: "batch", 2: "anchors"}  # shape(1, 116, 8400)
            dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"}  # shape(1,32,160,160)
        elif isinstance(self.model, DetectionModel):
            dynamic["output0"] = {0: "batch", 2: "anchors"}  # shape(1, 84, 8400)
        if self.args.nms:  # only batch size is dynamic with NMS
            dynamic["output0"].pop(2)
    if self.args.nms and self.model.task == "obb":
        self.args.opset = opset_version  # for NMSModel

    with arange_patch(self.args):
        export_onnx(
            NMSModel(self.model, self.args) if self.args.nms else self.model,
            self.im,
            f,
            opset=opset_version,
            input_names=["images"],
            output_names=output_names,
            dynamic=dynamic or None,
        )

    # Checks
    model_onnx = onnx.load(f)  # load onnx model

    # Simplify
    if self.args.simplify:
        try:
            import onnxslim

            LOGGER.info(f"{prefix} slimming with onnxslim {onnxslim.__version__}...")
            model_onnx = onnxslim.slim(model_onnx)

        except Exception as e:
            LOGGER.warning(f"{prefix} simplifier failure: {e}")

    # Metadata
    for k, v in self.metadata.items():
        meta = model_onnx.metadata_props.add()
        meta.key, meta.value = k, str(v)

    onnx.save(model_onnx, f)
    return f, model_onnx

@try_export
def export_openvino(self, prefix=colorstr("OpenVINO:")):
    """Export YOLO model to OpenVINO format."""
    # OpenVINO <= 2025.1.0 error on macOS 15.4+: https://github.com/openvinotoolkit/openvino/issues/30023"
    check_requirements("openvino>=2025.2.0" if MACOS and MACOS_VERSION >= "15.4" else "openvino>=2024.0.0")
    import openvino as ov

    LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...")
    assert TORCH_1_13, f"OpenVINO export requires torch>=1.13.0 but torch=={torch.__version__} is installed"
    ov_model = ov.convert_model(
        NMSModel(self.model, self.args) if self.args.nms else self.model,
        input=None if self.args.dynamic else [self.im.shape],
        example_input=self.im,
    )

    def serialize(ov_model, file):
        """Set RT info, serialize, and save metadata YAML."""
        ov_model.set_rt_info("YOLO", ["model_info", "model_type"])
        ov_model.set_rt_info(True, ["model_info", "reverse_input_channels"])
        ov_model.set_rt_info(114, ["model_info", "pad_value"])
        ov_model.set_rt_info([255.0], ["model_info", "scale_values"])
        ov_model.set_rt_info(self.args.iou, ["model_info", "iou_threshold"])
        ov_model.set_rt_info([v.replace(" ", "_") for v in self.model.names.values()], ["model_info", "labels"])
        if self.model.task != "classify":
            ov_model.set_rt_info("fit_to_window_letterbox", ["model_info", "resize_type"])

        ov.save_model(ov_model, file, compress_to_fp16=self.args.half)
        YAML.save(Path(file).parent / "metadata.yaml", self.metadata)  # add metadata.yaml

    if self.args.int8:
        fq = str(self.file).replace(self.file.suffix, f"_int8_openvino_model{os.sep}")
        fq_ov = str(Path(fq) / self.file.with_suffix(".xml").name)
        # INT8 requires nncf, nncf requires packaging>=23.2 https://github.com/openvinotoolkit/nncf/issues/3463
        check_requirements("packaging>=23.2")  # must be installed first to build nncf wheel
        check_requirements("nncf>=2.14.0")
        import nncf

        def transform_fn(data_item) -> np.ndarray:
            """Quantization transform function."""
            data_item: torch.Tensor = data_item["img"] if isinstance(data_item, dict) else data_item
            assert data_item.dtype == torch.uint8, "Input image must be uint8 for the quantization preprocessing"
            im = data_item.numpy().astype(np.float32) / 255.0  # uint8 to fp16/32 and 0-255 to 0.0-1.0
            return np.expand_dims(im, 0) if im.ndim == 3 else im

        # Generate calibration data for integer quantization
        ignored_scope = None
        if isinstance(self.model.model[-1], Detect):
            # Includes all Detect subclasses like Segment, Pose, OBB, WorldDetect, YOLOEDetect
            head_module_name = ".".join(list(self.model.named_modules())[-1][0].split(".")[:2])
            ignored_scope = nncf.IgnoredScope(  # ignore operations
                patterns=[
                    f".*{head_module_name}/.*/Add",
                    f".*{head_module_name}/.*/Sub*",
                    f".*{head_module_name}/.*/Mul*",
                    f".*{head_module_name}/.*/Div*",
                    f".*{head_module_name}\\.dfl.*",
                ],
                types=["Sigmoid"],
            )

        quantized_ov_model = nncf.quantize(
            model=ov_model,
            calibration_dataset=nncf.Dataset(self.get_int8_calibration_dataloader(prefix), transform_fn),
            preset=nncf.QuantizationPreset.MIXED,
            ignored_scope=ignored_scope,
        )
        serialize(quantized_ov_model, fq_ov)
        return fq, None

    f = str(self.file).replace(self.file.suffix, f"_openvino_model{os.sep}")
    f_ov = str(Path(f) / self.file.with_suffix(".xml").name)

    serialize(ov_model, f_ov)
    return f, None