add weather modes

add exports

classes.json

@@ -1,62 +1,16 @@
  [
-    {
-      "Id": 0,
-      "Name": "Armored-Vehicle",
-      "Color": "#80FF0000",
-      "ColorBrush": "#80FF0000"
-    },
-    {
-      "Id": 1,
-      "Name": "Truck",
-      "Color": "#8000FF00",
-      "ColorBrush": "#8000FF00"
-    },
-    {
-      "Id": 2,
-      "Name": "Vehicle",
-      "Color": "#800000FF",
-      "ColorBrush": "#800000FF"
-    },
-    {
-      "Id": 3,
-      "Name": "Artillery",
-      "Color": "#80FFFF00",
-      "ColorBrush": "#80FFFF00"
-    },
-    {
-      "Id": 4,
-      "Name": "Shadow",
-      "Color": "#80FF00FF",
-      "ColorBrush": "#80FF00FF"
-    },
-    {
-      "Id": 5,
-      "Name": "Trenches",
-      "Color": "#8000FFFF",
-      "ColorBrush": "#8000FFFF"
-    },
-    {
-      "Id": 6,
-      "Name": "Military-men",
-      "Color": "#80000000",
-      "ColorBrush": "#80000000"
-    },
-    {
-      "Id": 7,
-      "Name": "Tyre-tracks",
-      "Color": "#80800000",
-      "ColorBrush": "#80800000"
-    },
-    {
-      "Id": 8,
-      "Name": "Additional-armored-tank",
-      "Color": "#80008000",
-      "ColorBrush": "#80008000"
-    },
-    {
-      "Id": 9,
-      "Name": "Smoke",
-      "Color": "#80000080",
-      "ColorBrush": "#80000080"
-    }
+    { "Id": 0, "Name": "Armored-Vehicle", "Color": "#80FF0000" },
+    { "Id": 1, "Name": "Truck", "Color": "#8000FF00" },
+    { "Id": 2, "Name": "Vehicle", "Color": "#800000FF" },
+    { "Id": 3, "Name": "Artillery", "Color": "#80FFFF00" },
+    { "Id": 4, "Name": "Shadow", "Color": "#80FF00FF" },
+    { "Id": 5, "Name": "Trenches", "Color": "#8000FFFF" },
+    { "Id": 6, "Name": "Military-men", "Color": "#80000000" },
+    { "Id": 7, "Name": "Tyre-tracks", "Color": "#80800000" },
+    { "Id": 8, "Name": "Additional-armored-tank", "Color": "#80008000" },
+    { "Id": 9, "Name": "Smoke", "Color": "#80000080" },
+    { "Id": 10, "Name": "Plane", "Color": "#80000080" },
+    { "Id": 11, "Name": "Moto", "Color": "#80808000" },
+    { "Id": 12, "Name": "Camouflage-net", "Color": "#80800080" },
+    { "Id": 13, "Name": "Camouflage-branches", "Color": "#80008080" }
   ]

dto/annotationClass.py

@@ -1,7 +1,14 @@
 import json
+from enum import Enum
 from os.path import dirname, join
 
 
+class WeatherMode(Enum):
+    Norm = 0
+    Wint = 20
+    Night = 40
+
+
 class AnnotationClass:
     def __init__(self, id, name, color):
         self.id = id
@@ -13,7 +20,13 @@ class AnnotationClass:
         classes_path = join(dirname(dirname(__file__)), 'classes.json')
         with open(classes_path, 'r', encoding='utf-8') as f:
             j = json.loads(f.read())
-            return {cl['Id']: AnnotationClass(id=cl['Id'], name=cl['Name'], color=cl['Color']) for cl in j}
+            annotations_dict = {}
+            for mode in WeatherMode:
+                for cl in j:
+                    id = mode.value + cl['Id']
+                    name = cl['Name'] if mode.value == 0 else f'{cl['Name']}({mode.name})'
+                    annotations_dict[id] = AnnotationClass(id, name, cl['Color'])
+            return annotations_dict
 
     @property
     def color_tuple(self):

exports/export.py

@@ -0,0 +1,34 @@
+import shutil
+from os import path
+from pathlib import Path
+import netron
+from ultralytics import YOLO
+
+
+def export_rknn(model_path):
+    # model_onnx = export_onnx(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 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')

train.py

@@ -119,11 +119,11 @@ def check_label(label_path):
 def create_yaml():
     print('creating yaml...')
     lines = ['names:']
-    for c in annotation_classes:
-        lines.append(f'- {annotation_classes[c].name}')
-    classes_count = len(annotation_classes)
-    for c in range(DEFAULT_CLASS_NUM - classes_count):
-        lines.append(f'- Class-{c + classes_count + 1}')
+    for i in range(DEFAULT_CLASS_NUM):
+        if i in annotation_classes:
+            lines.append(f'- {annotation_classes[i].name}')
+        else:
+            lines.append(f'- Class-{i + 1}')
     lines.append(f'nc: {DEFAULT_CLASS_NUM}')
 
     lines.append(f'test: test/images')
@@ -175,7 +175,8 @@ def train_dataset(existing_date=None, from_scratch=False):
         cur_folder = f'{prefix}{existing_date}'
         cur_dataset = path.join(datasets_dir, f'{prefix}{existing_date}')
     else:
-        form_dataset(latest_date)
+        # form_dataset(latest_date)
+        # create_yaml()
         cur_folder = today_folder
         cur_dataset = today_dataset
 
@@ -186,7 +187,7 @@ def train_dataset(existing_date=None, from_scratch=False):
     yaml = abspath(path.join(cur_dataset, 'data.yaml'))
     results = model.train(data=yaml,
                            epochs=120,
-                           batch=14,
+                           batch=11,
                            imgsz=1280,
                            save_period=1,
                            workers=24)