[WIP] Model interface (#68)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Laughing-q <1185102784@qq.com>

ultralytics/yolo/v8/classify/train.py

@@ -13,8 +13,10 @@ class ClassificationTrainer(BaseTrainer):
     def set_model_attributes(self):
         self.model.names = self.data["names"]
 
-    def load_model(self, model_cfg, weights):
+    def load_model(self, model_cfg=None, weights=None):
         # TODO: why treat clf models as unique. We should have clf yamls?
+        if isinstance(weights, dict):  # yolo ckpt
+            weights = weights["model"]
         if weights and not weights.__class__.__name__.startswith("yolo"):  # torchvision
             model = weights
         else:

ultralytics/yolo/v8/detect/train.py

@@ -15,7 +15,7 @@ from .val import DetectionValidator
 # BaseTrainer python usage
 class DetectionTrainer(SegmentationTrainer):
 
-    def load_model(self, model_cfg, weights):
+    def load_model(self, model_cfg=None, weights=None):
         model = DetectionModel(model_cfg or weights["model"].yaml,
                                ch=3,
                                nc=self.data["nc"],

ultralytics/yolo/v8/segment/train.py

@@ -26,7 +26,7 @@ class SegmentationTrainer(BaseTrainer):
         batch["img"] = batch["img"].to(self.device, non_blocking=True).float() / 255
         return batch
 
-    def load_model(self, model_cfg, weights):
+    def load_model(self, model_cfg=None, weights=None):
         model = SegmentationModel(model_cfg or weights["model"].yaml,
                                   ch=3,
                                   nc=self.data["nc"],