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:
            msg = "Model is already in PyTorch format." if fmt == "pt" else f"Invalid export format='{fmt}'."
            raise ValueError(f"{msg} 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.export.imx import FXModel

        model = FXModel(model, self.imgsz)
    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

    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 TORCH_1_11, "CoreML export requires torch>=1.11"
    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

@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

@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:
            cuda_version = torch.version.cuda.split(".")[0]
            check_requirements(f"tensorrt-cu{cuda_version}>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
    onnx2engine(
        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

@try_export
def export_imx(self, prefix=colorstr("IMX:")):
    """Export YOLO model to IMX format."""
    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

    # 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)

    return torch2imx(
        self.model,
        self.file,
        self.args.conf,
        self.args.iou,
        self.args.max_det,
        metadata=self.metadata,
        dataset=self.get_int8_calibration_dataloader(prefix),
        prefix=prefix,
    )

@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

@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)

@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.67", "onnxruntime" + ("-gpu" if torch.cuda.is_available() else "")]
    check_requirements(requirements)
    import onnx  # noqa

    opset = self.args.opset or best_onnx_opset(onnx)
    LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__} opset {opset}...")
    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  # for NMSModel

    with arange_patch(self.args):
        torch2onnx(
            NMSModel(self.model, self.args) if self.args.nms else self.model,
            self.im,
            f,
            opset=opset,
            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)

    # IR version
    if getattr(model_onnx, "ir_version", 0) > 10:
        LOGGER.info(f"{prefix} limiting IR version {model_onnx.ir_version} to 10 for ONNXRuntime compatibility...")
        model_onnx.ir_version = 10

    onnx.save(model_onnx, f)
    return f

@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_2_1, f"OpenVINO export requires torch>=2.1 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

    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

@try_export
def export_paddle(self, prefix=colorstr("PaddlePaddle:")):
    """Export YOLO model to PaddlePaddle format."""
    assert not IS_JETSON, "Jetson Paddle exports not supported yet"
    check_requirements(
        (
            "paddlepaddle-gpu"
            if torch.cuda.is_available()
            else "paddlepaddle==3.0.0"  # pin 3.0.0 for ARM64
            if ARM64
            else "paddlepaddle>=3.0.0",
            "x2paddle",
        )
    )
    import x2paddle  # noqa
    from x2paddle.convert import pytorch2paddle  # noqa

    LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...")
    f = str(self.file).replace(self.file.suffix, f"_paddle_model{os.sep}")

    pytorch2paddle(module=self.model, save_dir=f, jit_type="trace", input_examples=[self.im])  # export
    YAML.save(Path(f) / "metadata.yaml", self.metadata)  # add metadata.yaml
    return f

@try_export
def export_pb(self, keras_model, prefix=colorstr("TensorFlow GraphDef:")):
    """Export YOLO model to TensorFlow GraphDef *.pb format https://github.com/leimao/Frozen-Graph-TensorFlow."""
    import tensorflow as tf  # noqa
    from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2  # noqa

    LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
    f = self.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

@try_export
def export_rknn(self, prefix=colorstr("RKNN:")):
    """Export YOLO model to RKNN format."""
    LOGGER.info(f"\n{prefix} starting export with rknn-toolkit2...")

    check_requirements("rknn-toolkit2")
    if IS_COLAB:
        # Prevent 'exit' from closing the notebook https://github.com/airockchip/rknn-toolkit2/issues/259
        import builtins

        builtins.exit = lambda: None

    from rknn.api import RKNN

    f = self.export_onnx()
    export_path = Path(f"{Path(f).stem}_rknn_model")
    export_path.mkdir(exist_ok=True)

    rknn = RKNN(verbose=False)
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=self.args.name)
    rknn.load_onnx(model=f)
    rknn.build(do_quantization=False)  # TODO: Add quantization support
    f = f.replace(".onnx", f"-{self.args.name}.rknn")
    rknn.export_rknn(f"{export_path / f}")
    YAML.save(export_path / "metadata.yaml", self.metadata)
    return export_path