Refactor constants management to use Pydantic BaseModel for configuration

- 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.
This commit is contained in:
Oleksandr Bezdieniezhnykh
2026-03-27 18:18:30 +02:00
parent b68c07b540
commit 142c6c4de8
106 changed files with 5706 additions and 654 deletions
+100 -24
View File
@@ -1,28 +1,105 @@
from os import path
azaion = '/azaion'
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-'
images = 'images'
labels = 'labels'
data_dir = path.join(azaion, 'data')
data_images_dir = path.join(data_dir, images)
data_labels_dir = path.join(data_dir, labels)
processed_dir = path.join(azaion, 'data-processed')
processed_images_dir = path.join(processed_dir, images)
processed_labels_dir = path.join(processed_dir, labels)
corrupted_dir = path.join(azaion, 'data-corrupted')
corrupted_images_dir = path.join(corrupted_dir, images)
corrupted_labels_dir = path.join(corrupted_dir, labels)
sample_dir = path.join(azaion, 'data-sample')
datasets_dir = path.join(azaion, 'datasets')
models_dir = path.join(azaion, 'models')
date_format = '%Y-%m-%d'
checkpoint_file = 'checkpoint.txt'
checkpoint_date_format = '%Y-%m-%d %H:%M:%S'
@@ -38,5 +115,4 @@ SMALL_SIZE_KB = 3
CDN_CONFIG = 'cdn.yaml'
MODELS_FOLDER = 'models'
CURRENT_PT_MODEL = path.join(models_dir, f'{prefix[:-1]}.pt')
CURRENT_ONNX_MODEL = path.join(models_dir, f'{prefix[:-1]}.onnx')
config: Config = Config.from_yaml(CONFIG_FILE)