# Copyright (c) 2018-2026, Texas Instruments # All Rights Reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: # # * Redistributions of source code must retain the above copyright notice, this # list of conditions and the following disclaimer. # # * Redistributions in binary form must reproduce the above copyright notice, # this list of conditions and the following disclaimer in the documentation # and/or other materials provided with the distribution. # # * Neither the name of the copyright holder nor the names of its # contributors may be used to endorse or promote products derived from # this software without specific prior written permission. # # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. # --------------------------------------------------------------------------- # Script: prepare_model.py # # Purpose: # Export one or more YOLOv8 model variants (n / s / m / l / x) to ONNX # (or any other Ultralytics-supported format) using the Ultralytics Python # API. Required packages (onnx, ultralytics) are installed automatically # at runtime if they are not already present. # # Usage examples: # # Export only yolov8n (default behaviour) # python prepare_model.py # # # Export specific variants # python prepare_model.py --models yolov8n yolov8s # # # Export all variants # python prepare_model.py --models yolov8n yolov8s yolov8m yolov8l yolov8x # # # Export to a custom directory in TorchScript format # python prepare_model.py --models yolov8n --format torchscript --output-dir ./exports # --------------------------------------------------------------------------- import argparse from pathlib import Path import sys import subprocess # --------------------------------------------------------------------------- # Helper utilities # --------------------------------------------------------------------------- def install_package(package: str) -> None: """Install a Python package at runtime using pip. This ensures that optional heavyweight dependencies (e.g. ``onnx``, ``ultralytics``) are available without requiring them to be listed in the project's top-level requirements file. Args: package (str): Package name as accepted by ``pip install`` (e.g. ``'onnx'``, ``'ultralytics>=8.0'``). """ subprocess.check_call([sys.executable, "-m", "pip", "install", package]) # --------------------------------------------------------------------------- # Core export logic # --------------------------------------------------------------------------- def export_model(checkpoint: str, output_format: str = 'onnx', output_dir: str = None, opset: int = None) -> str: """Export a YOLOv8 model checkpoint to the requested format. The function loads the model via the Ultralytics ``YOLO`` class and calls its ``export()`` method. When *checkpoint* is a bare filename such as ``'yolov8n.pt'``, Ultralytics will automatically download the pre-trained weights from the official release page on first use. Supported export formats (non-exhaustive): ``onnx``, ``torchscript``, ``tflite``, ``pb``, ``saved_model``, ``coreml``, ``paddle``, ``ncnn``. Args: checkpoint (str): Path to, or name of, the YOLO checkpoint file (e.g. ``'yolov8n.pt'`` or ``'/path/to/custom_model.pt'``). output_format (str): Target export format understood by Ultralytics. Defaults to ``'onnx'``. output_dir (str | None): Directory where the exported artefact will be written. The directory is created automatically if it does not exist. When ``None`` the Ultralytics default location (same folder as the source ``.pt`` file) is used. opset (int | None): ONNX opset version to use for export. Only applies when output_format is 'onnx'. When ``None``, uses Ultralytics default. Returns: str: Absolute path to the exported model file as reported by Ultralytics. """ # Import here so that the module can be imported without ultralytics # being installed (the install happens in __main__ before any export call). from ultralytics.models import YOLO print(f"\n[export] Loading checkpoint: {checkpoint}") model = YOLO(checkpoint) # Build keyword arguments for model.export() export_kwargs = { 'format': output_format, } # import torch # dummy_input = torch.zeros(3, 640, 640, dtype=torch.float32) # torch.onnx.export( # model= # ) if output_dir: # Resolve and create the output directory tree as needed output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) export_kwargs['project'] = str(output_dir) if opset is not None and output_format == 'onnx': export_kwargs['opset'] = opset print(f"[export] Exporting to format='{output_format}'" + (f", output_dir='{output_dir}'" if output_dir else "") + " …") exported_path = model.export(**export_kwargs) print(f"[export] Model exported successfully to: {exported_path}") return exported_path # --------------------------------------------------------------------------- # CLI argument parsing # --------------------------------------------------------------------------- # All recognised YOLOv8 size variants in ascending order of complexity ALL_MODELS = ['yolov8n', 'yolov8s', 'yolov8m', 'yolov8l', 'yolov8x'] # Default: export only the nano variant (smallest / fastest) DEFAULT_MODELS = ['yolov8n'] def parse_args() -> argparse.Namespace: """Parse command-line arguments for the export script. Returns: argparse.Namespace: Parsed argument object with the following fields: * ``models`` – list of model base-names to export (without ``.pt``) * ``format`` – Ultralytics export format string * ``output_dir`` – destination directory (``None`` → Ultralytics default) """ parser = argparse.ArgumentParser( description=( "Export YOLOv8 model variants to ONNX (or another format) " "using the Ultralytics Python API." ), formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument( '--models', nargs='+', choices=ALL_MODELS + ['all'], default=DEFAULT_MODELS, metavar='MODEL', help=( "One or more YOLOv8 variants to export. " f"Choices: {ALL_MODELS} or 'all'. " f"Default: {DEFAULT_MODELS}." ), ) parser.add_argument( '--list-models', action='store_true', help='List all supported YOLOv8 model variants and exit.', ) parser.add_argument( '--format', default='onnx', metavar='FORMAT', help=( "Ultralytics export format. " "Common values: onnx, torchscript, tflite, pb, saved_model, coreml." ), ) parser.add_argument( '--output-dir', default=None, metavar='DIR', help=( "Directory to save exported models. " "Created automatically if it does not exist. " "Defaults to the Ultralytics standard location when omitted." ), ) parser.add_argument( '--opset', type=int, default=19, metavar='VERSION', help=( "ONNX opset version to use for export (e.g., 19). " "Only applies when --format is 'onnx'. " "Defaults to Ultralytics default opset when omitted." ), ) return parser.parse_args() # --------------------------------------------------------------------------- # Entry point # --------------------------------------------------------------------------- if __name__ == '__main__': # ------------------------------------------------------------------ # 1. Parse CLI arguments # ------------------------------------------------------------------ args = parse_args() if args.list_models: print("Supported YOLOv8 model variants:") for m in ALL_MODELS: print(f" {m}") sys.exit(0) if 'all' in args.models: args.models = list(ALL_MODELS) # ------------------------------------------------------------------ # 2. Ensure runtime dependencies are available # ------------------------------------------------------------------ # 'onnx' is required for ONNX export; 'ultralytics' provides the YOLO API; # 'onnxslim' is used internally by ultralytics' export(simplify=True) (the # default), which replaced the older 'onnxsim' package; 'onnxruntime' is # used internally to validate/load the exported ONNX model. Installing # these upfront avoids ultralytics' own AutoUpdate, which fails offline. install_package('onnx') install_package('ultralytics') install_package('onnxslim') install_package('onnxruntime') print(f"Models to export : {args.models}") print(f"Export format : {args.format}") print(f"Output directory : {args.output_dir or '(Ultralytics default)'}") # ------------------------------------------------------------------ # 3. Export each requested model variant # ------------------------------------------------------------------ for model_name in args.models: # Append the '.pt' extension expected by Ultralytics checkpoint = f"{model_name}.pt" export_model( checkpoint=checkpoint, output_format=args.format, output_dir=args.output_dir, opset=args.opset, ) print("\n[done] All requested models exported.")