Commit 37aef137 authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): bug fix

parent a70cd4c7
Loading
Loading
Loading
Loading
+6 −7
Original line number Diff line number Diff line
import os.path

from hbutils.system import copy
from hbutils.testing import isolated_directory
from ultralytics import YOLO


def export_yolo_to_onnx(yolo: YOLO, onnx_filename, opset_version: int = 14,
                        no_optimize: bool = False):
    _current_path = os.path.abspath(os.curdir)
    with isolated_directory():
    if os.path.dirname(onnx_filename):
        os.makedirs(os.path.dirname(onnx_filename), exist_ok=True)
    copy(
        yolo.export(format='onnx', dynamic=True, simplify=not no_optimize, opset=opset_version),
            os.path.join(_current_path, onnx_filename)
        onnx_filename
    )