From ee370fec051c2acfb8facf3895b8b6343cd41efe Mon Sep 17 00:00:00 2001
From: Alex Bezdieniezhnykh <zxsanny@gmail.com>
Date: Wed, 12 Mar 2025 01:11:35 +0200
Subject: [PATCH] add conversion to rockchip

---
 exports/rockchip.py | 7 +++++++
 train.py            | 2 ++
 2 files changed, 9 insertions(+)
 create mode 100644 exports/rockchip.py

diff --git a/exports/rockchip.py b/exports/rockchip.py
new file mode 100644
index 0000000..a030a11
--- /dev/null
+++ b/exports/rockchip.py
@@ -0,0 +1,7 @@
+from ultralytics import YOLO
+
+
+def export_rknn(model_path):
+    model = YOLO(model_path)
+    model.export(format="rknn", name="rk3588")
+    pass
\ No newline at end of file
diff --git a/train.py b/train.py
index 05cd5f3..3bd4673 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.rockchip import export_rknn
 from security import Security
 from utils import Dotdict
 
@@ -207,6 +208,7 @@ def convert2rknn():
     model.export(format="onnx")
     pass
 
+
 def form_data_sample(size=300):
     images = []
     with scandir(processed_images_dir) as imd: