add export to tensorrt

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",
+    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_tensorrt('azaion-2025-03-10.pt', path.join(datasets_dir, 'azaion-2025-03-10', 'data.yaml'))
-    # export_rknn('azaion_2025-03-10.pt')
+    # export_rknn('azaion-2025-03-10.pt')
+    # export_onnx('azaion-2025-03-10.pt')

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()