azaion/ai-training
1
0
Fork 0
mirror of https://github.com/azaion/ai-training.git synced 2026-04-22 08:46:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
zxsanny
2025-02-21 14:26:54 +02:00
parent 6e4b0f40ef
commit ceb50bf48a
11 changed files with 205 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+62 -18
View File
@@ -1,10 +1,10 @@
import os
import random
import subprocess
from os import path, replace, remove, listdir, makedirs, scandir
from os.path import abspath
import shutil
import subprocess
from datetime import datetime
from os import path, replace, listdir, makedirs, scandir
from os.path import abspath
from pathlib import Path
from ultralytics import YOLO
from constants import (processed_images_dir,
@@ -12,13 +12,14 @@ from constants import (processed_images_dir,
annotation_classes,
prefix, date_format,
datasets_dir, models_dir,
corrupted_images_dir, corrupted_labels_dir)
corrupted_images_dir, corrupted_labels_dir, sample_dir)
today_folder = f'{prefix}{datetime.now():{date_format}}'
today_dataset = path.join(datasets_dir, today_folder)
train_set = 70
valid_set = 20
test_set = 10
old_images_percentage = 75
DEFAULT_CLASS_NUM = 80
@@ -26,32 +27,33 @@ DEFAULT_CLASS_NUM = 80
def form_dataset(from_date: datetime):
makedirs(today_dataset, exist_ok=True)
images = []
old_images = []
with scandir(processed_images_dir) as imd:
for image_file in imd:
if not image_file.is_file():
continue
mod_time = datetime.fromtimestamp(image_file.stat().st_mtime)
mod_time = datetime.fromtimestamp(image_file.stat().st_mtime).replace(hour=0, minute=0, second=0, microsecond=0)
if from_date is None:
images.append(image_file)
elif mod_time > from_date:
images.append(image_file)
else: # gather old images as well in order to avoid overfitting on the only new data.
old_images.append(image_file)
random.shuffle(old_images)
old_images_size = int(len(old_images) * old_images_percentage / 100.0)
print(f'Got {len(images)} new images and {old_images_size} of old images (to prevent overfitting). Shuffling them...')
images.extend(old_images[:old_images_size])
print('shuffling images')
random.shuffle(images)
train_size = int(len(images) * train_set / 100.0)
valid_size = int(len(images) * valid_set / 100.0)
print(f'copy train dataset, size: {train_size} annotations')
copy_annotations(images[:train_size], 'train')
print(f'copy valid set, size: {valid_size} annotations')
copy_annotations(images[train_size:train_size + valid_size], 'valid')
print(f'copy test set, size: {len(images) - train_size - valid_size} annotations')
copy_annotations(images[train_size + valid_size:], 'test')
print('creating yaml...')
create_yaml()
@@ -65,6 +67,8 @@ def copy_annotations(images, folder):
makedirs(corrupted_images_dir, exist_ok=True)
makedirs(corrupted_labels_dir, exist_ok=True)
copied = 0
print(f'Copying annotations to {destination_images} and {destination_labels} folders:')
for image in images:
label_name = f'{Path(image.path).stem}.txt'
label_path = path.join(processed_labels_dir, label_name)
@@ -75,6 +79,10 @@ def copy_annotations(images, folder):
shutil.copy(image.path, path.join(corrupted_images_dir, image.name))
shutil.copy(label_path, path.join(corrupted_labels_dir, label_name))
print(f'Label {label_path} is corrupted! Copy with its image to the corrupted directory ({corrupted_labels_dir})')
copied = copied + 1
if copied % 1000 == 0:
print(f'{copied} copied...')
print(f'Copied all {copied} annotations to {destination_images} and {destination_labels} folders')
def check_label(label_path):
@@ -90,6 +98,7 @@ 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}')
@@ -136,8 +145,7 @@ def get_latest_model():
last_model = sorted_dates[-1]
return last_model['date'], last_model['path']
def train_dataset(existing_date=None):
def train_dataset(existing_date=None, from_scratch=False):
latest_date, latest_model = get_latest_model()
if existing_date is not None:
@@ -148,12 +156,17 @@ def train_dataset(existing_date=None):
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'
model_name = latest_model if latest_model is not None and path.isfile(latest_model) and not from_scratch else 'yolov8m.yaml'
print(f'Initial model: {model_name}')
model = YOLO(model_name)
yaml = abspath(path.join(cur_dataset, 'data.yaml'))
results = model.train(data=yaml, epochs=100, batch=57, imgsz=640, save_period=1)
results = model.train(data=yaml,
epochs=120,
batch=14,
imgsz=1280,
save_period=1,
workers=24)
model_dir = path.join(models_dir, cur_folder)
shutil.copytree(results.save_dir, model_dir)
@@ -164,9 +177,40 @@ def train_dataset(existing_date=None):
def convert2rknn():
subprocess.call(['bash', 'convert.sh'], cwd="./orangepi5")
latest_date, latest_model = get_latest_model()
model = YOLO(latest_model)
model.export(format="onnx")
pass
def form_data_sample(size=300):
images = []
with scandir(processed_images_dir) as imd:
for image_file in imd:
if not image_file.is_file():
continue
images.append(image_file)
print('shuffling images')
random.shuffle(images)
images = images[:size]
shutil.rmtree(sample_dir, ignore_errors=True)
makedirs(sample_dir, exist_ok=True)
lines = []
for image in images:
shutil.copy(image.path, path.join(sample_dir, image.name))
lines.append(f'./{image.name}')
with open(path.join(sample_dir, 'azaion_subset.txt'), 'w', encoding='utf-8') as f:
f.writelines([f'{line}\n' for line in lines])
def validate(model_path):
model = YOLO(model_path)
metrics = model.val()
pass
if __name__ == '__main__':
train_dataset()
train_dataset('2024-10-26', from_scratch=True)
validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
form_data_sample(500)
convert2rknn()