mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 11:26:36 +00:00
add weather modes
add exports
This commit is contained in:
+14
-60
@@ -1,62 +1,16 @@
|
|||||||
[
|
[
|
||||||
{
|
{ "Id": 0, "Name": "Armored-Vehicle", "Color": "#80FF0000" },
|
||||||
"Id": 0,
|
{ "Id": 1, "Name": "Truck", "Color": "#8000FF00" },
|
||||||
"Name": "Armored-Vehicle",
|
{ "Id": 2, "Name": "Vehicle", "Color": "#800000FF" },
|
||||||
"Color": "#80FF0000",
|
{ "Id": 3, "Name": "Artillery", "Color": "#80FFFF00" },
|
||||||
"ColorBrush": "#80FF0000"
|
{ "Id": 4, "Name": "Shadow", "Color": "#80FF00FF" },
|
||||||
},
|
{ "Id": 5, "Name": "Trenches", "Color": "#8000FFFF" },
|
||||||
{
|
{ "Id": 6, "Name": "Military-men", "Color": "#80000000" },
|
||||||
"Id": 1,
|
{ "Id": 7, "Name": "Tyre-tracks", "Color": "#80800000" },
|
||||||
"Name": "Truck",
|
{ "Id": 8, "Name": "Additional-armored-tank", "Color": "#80008000" },
|
||||||
"Color": "#8000FF00",
|
{ "Id": 9, "Name": "Smoke", "Color": "#80000080" },
|
||||||
"ColorBrush": "#8000FF00"
|
{ "Id": 10, "Name": "Plane", "Color": "#80000080" },
|
||||||
},
|
{ "Id": 11, "Name": "Moto", "Color": "#80808000" },
|
||||||
{
|
{ "Id": 12, "Name": "Camouflage-net", "Color": "#80800080" },
|
||||||
"Id": 2,
|
{ "Id": 13, "Name": "Camouflage-branches", "Color": "#80008080" }
|
||||||
"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"
|
|
||||||
}
|
|
||||||
]
|
]
|
||||||
|
|||||||
+14
-1
@@ -1,7 +1,14 @@
|
|||||||
import json
|
import json
|
||||||
|
from enum import Enum
|
||||||
from os.path import dirname, join
|
from os.path import dirname, join
|
||||||
|
|
||||||
|
|
||||||
|
class WeatherMode(Enum):
|
||||||
|
Norm = 0
|
||||||
|
Wint = 20
|
||||||
|
Night = 40
|
||||||
|
|
||||||
|
|
||||||
class AnnotationClass:
|
class AnnotationClass:
|
||||||
def __init__(self, id, name, color):
|
def __init__(self, id, name, color):
|
||||||
self.id = id
|
self.id = id
|
||||||
@@ -13,7 +20,13 @@ class AnnotationClass:
|
|||||||
classes_path = join(dirname(dirname(__file__)), 'classes.json')
|
classes_path = join(dirname(dirname(__file__)), 'classes.json')
|
||||||
with open(classes_path, 'r', encoding='utf-8') as f:
|
with open(classes_path, 'r', encoding='utf-8') as f:
|
||||||
j = json.loads(f.read())
|
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
|
@property
|
||||||
def color_tuple(self):
|
def color_tuple(self):
|
||||||
|
|||||||
@@ -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')
|
||||||
@@ -119,11 +119,11 @@ def check_label(label_path):
|
|||||||
def create_yaml():
|
def create_yaml():
|
||||||
print('creating yaml...')
|
print('creating yaml...')
|
||||||
lines = ['names:']
|
lines = ['names:']
|
||||||
for c in annotation_classes:
|
for i in range(DEFAULT_CLASS_NUM):
|
||||||
lines.append(f'- {annotation_classes[c].name}')
|
if i in annotation_classes:
|
||||||
classes_count = len(annotation_classes)
|
lines.append(f'- {annotation_classes[i].name}')
|
||||||
for c in range(DEFAULT_CLASS_NUM - classes_count):
|
else:
|
||||||
lines.append(f'- Class-{c + classes_count + 1}')
|
lines.append(f'- Class-{i + 1}')
|
||||||
lines.append(f'nc: {DEFAULT_CLASS_NUM}')
|
lines.append(f'nc: {DEFAULT_CLASS_NUM}')
|
||||||
|
|
||||||
lines.append(f'test: test/images')
|
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_folder = f'{prefix}{existing_date}'
|
||||||
cur_dataset = path.join(datasets_dir, f'{prefix}{existing_date}')
|
cur_dataset = path.join(datasets_dir, f'{prefix}{existing_date}')
|
||||||
else:
|
else:
|
||||||
form_dataset(latest_date)
|
# form_dataset(latest_date)
|
||||||
|
# create_yaml()
|
||||||
cur_folder = today_folder
|
cur_folder = today_folder
|
||||||
cur_dataset = today_dataset
|
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'))
|
yaml = abspath(path.join(cur_dataset, 'data.yaml'))
|
||||||
results = model.train(data=yaml,
|
results = model.train(data=yaml,
|
||||||
epochs=120,
|
epochs=120,
|
||||||
batch=14,
|
batch=11,
|
||||||
imgsz=1280,
|
imgsz=1280,
|
||||||
save_period=1,
|
save_period=1,
|
||||||
workers=24)
|
workers=24)
|
||||||
|
|||||||
Reference in New Issue
Block a user