From 3113d59a3ad9fb9fdd3f8458e9a79d14b312abaa Mon Sep 17 00:00:00 2001 From: zxsanny Date: Thu, 13 Mar 2025 16:39:57 +0200 Subject: [PATCH] add weather modes add exports --- classes.json | 74 ++++++++---------------------------------- dto/annotationClass.py | 15 ++++++++- exports/export.py | 34 +++++++++++++++++++ train.py | 15 +++++---- 4 files changed, 70 insertions(+), 68 deletions(-) create mode 100644 exports/export.py diff --git a/classes.json b/classes.json index f127ffb..a83dd9e 100644 --- a/classes.json +++ b/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" } ] diff --git a/dto/annotationClass.py b/dto/annotationClass.py index ef8938c..558d6fa 100644 --- a/dto/annotationClass.py +++ b/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): diff --git a/exports/export.py b/exports/export.py new file mode 100644 index 0000000..9f6e305 --- /dev/null +++ b/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') \ No newline at end of file diff --git a/train.py b/train.py index 02424dc..4f5b992 100644 --- a/train.py +++ b/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)