mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 11:06:35 +00:00
zxsanny
81 lines
2.3 KiB
Python
81 lines
2.3 KiB
Python
import shutil
|
|
from os import path, scandir, makedirs
|
|
from pathlib import Path
|
|
import random
|
|
|
|
import netron
|
|
import yaml
|
|
from ultralytics import YOLO
|
|
|
|
from constants import datasets_dir, processed_images_dir
|
|
|
|
|
|
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 export_tensorrt(model_path, dataset_yaml):
|
|
form_data_sample(path.join(path.dirname(dataset_yaml), 'minival', 'images'))
|
|
model = YOLO(model_path)
|
|
with open(dataset_yaml, 'r') as file:
|
|
yaml_data = yaml.safe_load(file) or {}
|
|
yaml_data['minival'] = 'minival/images'
|
|
with open(dataset_yaml, 'w') as file:
|
|
yaml.dump(yaml_data, file)
|
|
|
|
model.export(
|
|
format='engine',
|
|
batch=4,
|
|
half=True,
|
|
nms=True,
|
|
data=dataset_yaml,
|
|
split='minival'
|
|
)
|
|
|
|
def form_data_sample(destination_path, size=500, write_txt_log=False):
|
|
images = []
|
|
with scandir(processed_images_dir) as imd:
|
|
for image_file in imd:
|
|
if not image_file.is_file():
|
|
continue
|
|
images.append(image_file)
|
|
print('shuffling images')
|
|
random.shuffle(images)
|
|
images = images[:size]
|
|
|
|
shutil.rmtree(destination_path, ignore_errors=True)
|
|
makedirs(destination_path, exist_ok=True)
|
|
|
|
lines = []
|
|
for image in images:
|
|
shutil.copy(image.path, path.join(destination_path, image.name))
|
|
lines.append(f'./{image.name}')
|
|
if write_txt_log:
|
|
with open(path.join(destination_path, 'azaion_subset.txt'), 'w', encoding='utf-8') as f:
|
|
f.writelines([f'{line}\n' for line in lines])
|
|
|
|
def show_model(model: str = None):
|
|
netron.start(model)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
export_tensorrt('azaion-2025-03-10.pt', path.join(datasets_dir, 'azaion-2025-03-10', 'data.yaml'))
|
|
# export_rknn('azaion-2025-03-10.pt')
|
|
# export_onnx('azaion-2025-03-10.pt') |