import shutil
from os import path
from pathlib import Path
import netron
from ultralytics import YOLO


def export_rknn(model_path):
    model = YOLO(model_path)
    model.export(format="rknn", name="rk3588", simplify=True)
    model_stem = Path(model_path).stem
    folder_name = f'{model_stem}_rknn_model'
    shutil.move(path.join(folder_name, f'{Path(model_path).stem}-rk3588.rknn'), f'{model_stem}.rknn')
    shutil.rmtree(folder_name)
    pass


def export_onnx(model_path):
    model = YOLO(model_path)
    model.export(format="onnx",
                 imgsz=1280,
                 batch=2,
                 simplify=True,
                 nms=True)
    return Path(model_path).stem + '.onnx'


def show_model(model: str = None):
    netron.start(model)


if __name__ == '__main__':
    export_onnx('azaion-2024-10-26.pt')
    show_model('azaion-2024-10-26.onnx')
    # export_rknn('azaion_2025-03-10.pt')