Model interface enhancement (#106)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

ultralytics/yolo/engine/model.py

@@ -43,6 +43,7 @@ class YOLO:
         self.trainer = None
         self.task = None
         self.ckpt = None
+        self.overrides = {}
 
     def new(self, cfg: str):
         """
@@ -69,6 +70,10 @@ class YOLO:
         """
         self.ckpt = torch.load(weights, map_location="cpu")
         self.task = self.ckpt["train_args"]["task"]
+        self.overrides = dict(self.ckpt["train_args"])
+        self.overrides["device"] = ''  # reset device
+        LOGGER.info("Device has been reset to ''")
+
         self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._guess_ops_from_task(
             task=self.task)
         self.model = attempt_load_weights(weights)
@@ -107,6 +112,7 @@ class YOLO:
         source (str): Accepts all source types accepted by yolo
         **kwargs : Any other args accepted by the predictors. Too see all args check 'configuration' section in the docs
         """
+        kwargs.update(self.overrides)
         predictor = self.PredictorClass(overrides=kwargs)
 
         # check size type
@@ -119,7 +125,7 @@ class YOLO:
         predictor.setup(model=self.model, source=source)
         predictor()
 
-    def val(self, data, **kwargs):
+    def val(self, data=None, **kwargs):
         """
         Validate a model on a given dataset
 
@@ -130,8 +136,9 @@ class YOLO:
         if not self.model:
             raise Exception("model not initialized!")
 
+        kwargs.update(self.overrides)
         args = get_config(config=DEFAULT_CONFIG, overrides=kwargs)
-        args.data = data
+        args.data = data or args.data
         args.task = self.task
 
         validator = self.ValidatorClass(args=args)

ultralytics/yolo/engine/predictor.py

@@ -86,10 +86,15 @@ class BasePredictor:
 
         # data
         if self.data:
-            if self.data.endswith(".yaml"):
-                self.data = check_dataset_yaml(self.data)
-            else:
-                self.data = check_dataset(self.data)
+            try:
+                if self.data.endswith(".yaml"):
+                    self.data = check_dataset_yaml(self.data)
+                else:
+                    self.data = check_dataset(self.data)
+            except AssertionError as e:
+                LOGGER.info(f"Error ocurred: {e}")
+            finally:
+                LOGGER.info("Predictor will continue without reading the dataset")
 
         # model
         device = select_device(self.args.device)

ultralytics/yolo/engine/trainer.py

@@ -46,10 +46,15 @@ class BaseTrainer:
         self.validator = None
         self.model = None
         self.callbacks = defaultdict(list)
-        self.save_dir = increment_path(Path(self.args.project) / self.args.name, exist_ok=self.args.exist_ok)
+
+        # dirs
+        project = overrides.get("project") or self.args.task
+        name = overrides.get("name") or self.args.mode
+        self.save_dir = increment_path(Path("runs") / project / name, exist_ok=self.args.exist_ok)
         self.wdir = self.save_dir / 'weights'  # weights dir
         self.wdir.mkdir(parents=True, exist_ok=True)  # make dir
         self.last, self.best = self.wdir / 'last.pt', self.wdir / 'best.pt'  # checkpoint paths
+
         self.batch_size = self.args.batch_size
         self.epochs = self.args.epochs
         self.start_epoch = 0

ultralytics/yolo/utils/configs/__init__.py

@@ -6,7 +6,7 @@ from omegaconf import DictConfig, OmegaConf
 from ultralytics.yolo.utils.configs.hydra_patch import check_config_mismatch
 
 
-def get_config(config: Union[str, DictConfig], overrides: Union[str, Dict]):
+def get_config(config: Union[str, DictConfig], overrides: Union[str, Dict] = {}):
     """
     Accepts yaml file name or DictConfig containing experiment configuration.
     Returns training args namespace