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

fix folder names, refactoring

Browse Source 
move to yolov8m
use checkpoint.txt instead of yaml
This commit is contained in:
Alex Bezdieniezhnykh
2024-06-29 22:18:34 +03:00
parent 3fd72202f2
commit bb1dbfe1e7
10 changed files with 116 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+25 -23
View File
@@ -3,19 +3,20 @@ from os.path import abspath
import shutil
from datetime import datetime
from pathlib import Path
from ultralytics import YOLOv10
from constants import current_images_dir, current_labels_dir, annotation_classes, prefix, date_format
from ultralytics import YOLO
from constants import processed_images_dir, processed_labels_dir, annotation_classes, prefix, date_format, datasets_dir, models_dir
latest_model = f'models/{prefix}latest.pt'
latest_model = path.join(models_dir, f'{prefix}latest.pt')
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
def form_dataset():
makedirs(path.join('datasets', today_folder), exist_ok=True)
images = listdir(current_images_dir)
makedirs(today_dataset, exist_ok=True)
images = listdir(processed_images_dir)
train_size = int(len(images) * train_set / 100.0)
valid_size = int(len(images) * valid_set / 100.0)
@@ -28,15 +29,14 @@ def form_dataset():
def move_annotations(images, folder):
today_dataset = path.join('datasets', today_folder)
destination_images = path.join(today_dataset, folder, 'images')
makedirs(destination_images, exist_ok=True)
destination_labels = path.join(today_dataset, folder, 'labels')
makedirs(destination_labels, exist_ok=True)
for image_name in images:
image_path = path.join(current_images_dir, image_name)
image_path = path.join(processed_images_dir, image_name)
label_name = f'{Path(image_name).stem}.txt'
label_path = path.join(current_labels_dir, label_name)
label_path = path.join(processed_labels_dir, label_name)
if not check_label(label_path):
remove(image_path)
else:
@@ -77,38 +77,40 @@ def create_yaml():
lines.append(f'val: valid/images')
lines.append('')
today_yaml = abspath(path.join('datasets', today_folder, 'data.yaml'))
today_yaml = abspath(path.join(today_dataset, 'data.yaml'))
with open(today_yaml, 'w', encoding='utf-8') as f:
f.writelines([f'{line}\n' for line in lines])
def revert_to_current(date):
def revert_to_processed_data(date):
def revert_dir(src_dir, dest_dir):
for file in listdir(src_dir):
s = path.join(src_dir, file)
d = path.join(dest_dir, file)
replace(s, d)
date_dataset = path.join('datasets', f'{prefix}{date}')
current_dataset = path.join('datasets', f'{prefix}current')
date_dataset = path.join(datasets_dir, f'{prefix}{date}')
makedirs(processed_images_dir, exist_ok=True)
makedirs(processed_labels_dir, exist_ok=True)
for subset in ['test', 'train', 'valid']:
revert_dir(path.join(date_dataset, subset, 'images'), path.join(current_dataset, 'images'))
revert_dir(path.join(date_dataset, subset, 'labels'), path.join(current_dataset, 'labels'))
revert_dir(path.join(date_dataset, subset, 'images'), processed_images_dir)
revert_dir(path.join(date_dataset, subset, 'labels'), processed_labels_dir)
shutil.rmtree(date_dataset)
if __name__ == '__main__':
# form_dataset()
# create_yaml()
m = latest_model or 'yolov10x.yaml'
print(f'Initial model: {m}')
model = YOLOv10(latest_model or 'yolov10x.yaml')
model_name = latest_model if 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-06-18')
folder = f'{prefix}2024-06-18'
yaml = abspath(path.join('datasets', folder, 'data.yaml'))
results = model.train(data=yaml, epochs=100, batch=10, imgsz=640, save_period=1)
cur_folder = today_dataset
yaml = abspath(path.join(cur_folder, 'data.yaml'))
results = model.train(data=yaml, epochs=100, batch=55, imgsz=640, save_period=1)
shutil.copy(f'{results.save_dir}/weights/best.pt', latest_model)
shutil.copytree(results.save_dir, f'models/{folder}')
shutil.copytree(results.save_dir, path.join(models_dir, cur_folder))
shutil.rmtree('runs')
shutil.rmtree('models/zombobase-latest')
shutil.rmtree(path.join(models_dir, f'{prefix}latest'))