diff --git a/exports/export.py b/exports/export.py index 9f6e305..20c7f54 100644 --- a/exports/export.py +++ b/exports/export.py @@ -1,9 +1,14 @@ import shutil -from os import path +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_onnx = export_onnx(model_path) @@ -18,17 +23,60 @@ def export_rknn(model_path): def export_onnx(model_path): model = YOLO(model_path) - model.export(format="onnx", - imgsz=1280, - batch=2, - simplify=True, - nms=True) + 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__': - show_model('azaion_2025-03-10.rknn') - # export_rknn('azaion_2025-03-10.pt') \ No newline at end of file + 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') \ No newline at end of file diff --git a/train.py b/train.py index 4f5b992..6bacd07 100644 --- a/train.py +++ b/train.py @@ -20,6 +20,7 @@ from constants import (processed_images_dir, prefix, date_format, datasets_dir, models_dir, corrupted_images_dir, corrupted_labels_dir, sample_dir) +from exports.export import form_data_sample from security import Security from utils import Dotdict @@ -209,28 +210,6 @@ def convert2rknn(): pass -def form_data_sample(size=300): - 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(sample_dir, ignore_errors=True) - makedirs(sample_dir, exist_ok=True) - - lines = [] - for image in images: - shutil.copy(image.path, path.join(sample_dir, image.name)) - lines.append(f'./{image.name}') - with open(path.join(sample_dir, 'azaion_subset.txt'), 'w', encoding='utf-8') as f: - f.writelines([f'{line}\n' for line in lines]) - - def validate(model_path): model = YOLO(model_path) metrics = model.val()