mirror of
https://github.com/azaion/ai-training.git
synced 2026-04-22 11:26:36 +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:
+13
-5
@@ -11,7 +11,6 @@ from ultralytics import YOLO
|
||||
import constants
|
||||
from api_client import ApiClient, ApiCredentials
|
||||
from cdn_manager import CDNManager, CDNCredentials
|
||||
from constants import datasets_dir, processed_images_dir
|
||||
from security import Security
|
||||
from utils import Dotdict
|
||||
|
||||
@@ -26,7 +25,9 @@ def export_rknn(model_path):
|
||||
pass
|
||||
|
||||
|
||||
def export_onnx(model_path, batch_size=4):
|
||||
def export_onnx(model_path, batch_size=None):
|
||||
if batch_size is None:
|
||||
batch_size = constants.config.export.onnx_batch
|
||||
model = YOLO(model_path)
|
||||
onnx_path = Path(model_path).stem + '.onnx'
|
||||
if path.exists(onnx_path):
|
||||
@@ -34,11 +35,18 @@ def export_onnx(model_path, batch_size=4):
|
||||
|
||||
model.export(
|
||||
format="onnx",
|
||||
imgsz=1280,
|
||||
imgsz=constants.config.export.onnx_imgsz,
|
||||
batch=batch_size,
|
||||
simplify=True,
|
||||
nms=True,
|
||||
device=0
|
||||
)
|
||||
|
||||
|
||||
def export_coreml(model_path):
|
||||
model = YOLO(model_path)
|
||||
model.export(
|
||||
format="coreml",
|
||||
imgsz=constants.config.export.onnx_imgsz,
|
||||
)
|
||||
|
||||
|
||||
@@ -54,7 +62,7 @@ def export_tensorrt(model_path):
|
||||
|
||||
def form_data_sample(destination_path, size=500, write_txt_log=False):
|
||||
images = []
|
||||
with scandir(processed_images_dir) as imd:
|
||||
with scandir(constants.config.processed_images_dir) as imd:
|
||||
for image_file in imd:
|
||||
if not image_file.is_file():
|
||||
continue
|
||||
|
||||
Reference in New Issue
Block a user