mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 06:56:34 +00:00
Oleksandr Bezdieniezhnykh
142c6c4de8
- Replaced module-level path variables in constants.py with a structured Pydantic Config class. - Updated all relevant modules (train.py, augmentation.py, exports.py, dataset-visualiser.py, manual_run.py) to access paths through the new config structure. - Fixed bugs related to image processing and model saving. - Enhanced test infrastructure to accommodate the new configuration approach. This refactor improves code maintainability and clarity by centralizing configuration management.
119 lines
2.8 KiB
Python
119 lines
2.8 KiB
Python
from os import path
|
|
|
|
import yaml
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class DirsConfig(BaseModel):
|
|
root: str = '/azaion'
|
|
|
|
|
|
class TrainingConfig(BaseModel):
|
|
model: str = 'yolo11m.yaml'
|
|
epochs: int = 120
|
|
batch: int = 11
|
|
imgsz: int = 1280
|
|
save_period: int = 1
|
|
workers: int = 24
|
|
|
|
|
|
class ExportConfig(BaseModel):
|
|
onnx_imgsz: int = 1280
|
|
onnx_batch: int = 4
|
|
|
|
|
|
class Config(BaseModel):
|
|
dirs: DirsConfig = DirsConfig()
|
|
training: TrainingConfig = TrainingConfig()
|
|
export: ExportConfig = ExportConfig()
|
|
|
|
@property
|
|
def azaion(self) -> str:
|
|
return self.dirs.root
|
|
|
|
@property
|
|
def data_dir(self) -> str:
|
|
return path.join(self.dirs.root, 'data')
|
|
|
|
@property
|
|
def data_images_dir(self) -> str:
|
|
return path.join(self.data_dir, 'images')
|
|
|
|
@property
|
|
def data_labels_dir(self) -> str:
|
|
return path.join(self.data_dir, 'labels')
|
|
|
|
@property
|
|
def processed_dir(self) -> str:
|
|
return path.join(self.dirs.root, 'data-processed')
|
|
|
|
@property
|
|
def processed_images_dir(self) -> str:
|
|
return path.join(self.processed_dir, 'images')
|
|
|
|
@property
|
|
def processed_labels_dir(self) -> str:
|
|
return path.join(self.processed_dir, 'labels')
|
|
|
|
@property
|
|
def corrupted_dir(self) -> str:
|
|
return path.join(self.dirs.root, 'data-corrupted')
|
|
|
|
@property
|
|
def corrupted_images_dir(self) -> str:
|
|
return path.join(self.corrupted_dir, 'images')
|
|
|
|
@property
|
|
def corrupted_labels_dir(self) -> str:
|
|
return path.join(self.corrupted_dir, 'labels')
|
|
|
|
@property
|
|
def sample_dir(self) -> str:
|
|
return path.join(self.dirs.root, 'data-sample')
|
|
|
|
@property
|
|
def datasets_dir(self) -> str:
|
|
return path.join(self.dirs.root, 'datasets')
|
|
|
|
@property
|
|
def models_dir(self) -> str:
|
|
return path.join(self.dirs.root, 'models')
|
|
|
|
@property
|
|
def current_pt_model(self) -> str:
|
|
return path.join(self.models_dir, f'{prefix[:-1]}.pt')
|
|
|
|
@property
|
|
def current_onnx_model(self) -> str:
|
|
return path.join(self.models_dir, f'{prefix[:-1]}.onnx')
|
|
|
|
@classmethod
|
|
def from_yaml(cls, config_file: str, root: str = None) -> 'Config':
|
|
try:
|
|
with open(config_file) as f:
|
|
data = yaml.safe_load(f) or {}
|
|
except FileNotFoundError:
|
|
data = {}
|
|
if root is not None:
|
|
data.setdefault('dirs', {})['root'] = root
|
|
return cls(**data)
|
|
|
|
|
|
prefix = 'azaion-'
|
|
date_format = '%Y-%m-%d'
|
|
checkpoint_file = 'checkpoint.txt'
|
|
checkpoint_date_format = '%Y-%m-%d %H:%M:%S'
|
|
|
|
CONFIG_FILE = 'config.yaml'
|
|
|
|
JPG_EXT = '.jpg'
|
|
TXT_EXT = '.txt'
|
|
|
|
OFFSET_FILE = 'offset.yaml'
|
|
|
|
SMALL_SIZE_KB = 3
|
|
CDN_CONFIG = 'cdn.yaml'
|
|
MODELS_FOLDER = 'models'
|
|
|
|
config: Config = Config.from_yaml(CONFIG_FILE)
|