mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 11:16:35 +00:00
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:
+100
-24
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user