diff --git a/src/model/yolo_wrapper.py b/src/model/yolo_wrapper.py index db95a2c..d5e31b3 100644 --- a/src/model/yolo_wrapper.py +++ b/src/model/yolo_wrapper.py @@ -80,7 +80,8 @@ class YOLOWrapper: Dictionary with training results """ if self.model is None: - self.load_model() + if not self.load_model(): + raise RuntimeError(f"Failed to load model from {self.model_path}") try: logger.info(f"Starting training: {name}") @@ -122,7 +123,8 @@ class YOLOWrapper: Dictionary with validation metrics """ if self.model is None: - self.load_model() + if not self.load_model(): + raise RuntimeError(f"Failed to load model from {self.model_path}") try: logger.info(f"Starting validation on {split} split") @@ -163,7 +165,8 @@ class YOLOWrapper: List of detection dictionaries """ if self.model is None: - self.load_model() + if not self.load_model(): + raise RuntimeError(f"Failed to load model from {self.model_path}") prepared_source, cleanup_path = self._prepare_source(source) @@ -211,7 +214,8 @@ class YOLOWrapper: Path to exported model """ if self.model is None: - self.load_model() + if not self.load_model(): + raise RuntimeError(f"Failed to load model from {self.model_path}") try: logger.info(f"Exporting model to {format} format")