mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 12:36:36 +00:00
Alex Bezdieniezhnykh
35 lines
911 B
Python
35 lines
911 B
Python
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') |