add rknn conversion - install and use scripts, auto convert to rknn after AI training is done and put pt and rknn models to /azaion/models directory

train.py

@@ -1,4 +1,6 @@
+import os
 import random
+import subprocess
 from os import path, replace, remove, listdir, makedirs, scandir
 from os.path import abspath
 import shutil
@@ -127,7 +129,7 @@ def get_latest_model():
         dir_model_path = path.join(models_dir, d, 'weights', 'best.pt')
         return {'date': dir_date, 'path': dir_model_path}
 
-    dates = [convert(d) for d in listdir(models_dir)]
+    dates = [convert(d) for d in next(os.walk(models_dir))[1]]
     sorted_dates = list(sorted(dates, key=lambda x: x['date']))
     if len(sorted_dates) == 0:
         return None, None
@@ -135,21 +137,36 @@ def get_latest_model():
     return last_model['date'], last_model['path']
 
 
-if __name__ == '__main__':
+def train_dataset(existing_date=None):
     latest_date, latest_model = get_latest_model()
-    # create_yaml()
-    # form_dataset(latest_date)
 
+    if existing_date is not None:
+        cur_folder = f'{prefix}{existing_date}'
+        cur_dataset = path.join(datasets_dir, f'{prefix}{existing_date}')
+    else:
+        form_dataset(latest_date)
+        cur_folder = today_folder
+        cur_dataset = today_dataset
 
     model_name = latest_model if latest_model is not None and path.isfile(latest_model) else 'yolov8m.yaml'
     print(f'Initial model: {model_name}')
     model = YOLO(model_name)
-    
-    cur_folder = path.join(datasets_dir, f'{prefix}2024-09-19')
 
-    # cur_folder = today_dataset
-    yaml = abspath(path.join(cur_folder, 'data.yaml'))
-    results = model.train(data=yaml, epochs=100, batch=58, imgsz=640)
+    yaml = abspath(path.join(cur_dataset, 'data.yaml'))
+    results = model.train(data=yaml, epochs=100, batch=57, imgsz=640, save_period=1)
 
-    shutil.copytree(results.save_dir, path.join(models_dir, today_folder))
+    model_dir = path.join(models_dir, cur_folder)
+    shutil.copytree(results.save_dir, model_dir)
+
+    shutil.copy(path.join(model_dir, 'weights', 'best.pt'), path.join(models_dir, f'{prefix[:-1]}.pt'))
     shutil.rmtree('runs')
+
+
+def convert2rknn():
+    subprocess.call(['bash', 'convert.sh'], cwd="./orangepi5")
+    pass
+
+
+if __name__ == '__main__':
+    train_dataset()
+    convert2rknn()