Make YOLO a module (#111)

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>

README.md

@@ -16,19 +16,21 @@ pip install -e .
 ### 1. CLI
 To simply use the latest Ultralytics YOLO models
 ```bash
-yolo task=detect    mode=train  model=s.yaml ...
-          classify       infer        s-cls.yaml
-          segment        val          s-seg.yaml
+yolo task=detect    mode=train     model=yolov8n.yaml ...
+          classify       predict         yolov8n-cls.yaml
+          segment        val             yolov8n-seg.yaml
 ```
 ### 2. Python SDK
 To use pythonic interface of Ultralytics YOLO model
 ```python
-import ultralytics
 from ultralytics import YOLO
 
-model = YOLO()
-model.new("s-seg.yaml") # automatically detects task type
-model.load("s-seg.pt") # load checkpoint
-model.train(data="coco128-segments", epochs=1, lr0=0.01, ...)
+model = YOLO.new('yolov8n.yaml')  # create a new model from scratch
+model = YOLO.load('yolov8n.pt')  # load a pretrained model (recommended for best training results)
+
+results = model.train(data='coco128.yaml', epochs=100, imgsz=640, ...)
+results = model.val()
+results = model.predict(source='bus.jpg')
+success = model.export(format='onnx')
 ```
 If you're looking to modify YOLO for R&D or to build on top of it, refer to [Using Trainer]() Guide on our docs.

tests/test_model.py

@@ -3,45 +3,49 @@ import torch
 from ultralytics import YOLO
 
 
+def test_model_init():
+    model = YOLO.new("yolov8n.yaml")
+    model.info()
+    try:
+        YOLO()
+    except Exception:
+        print("Successfully caught constructor assert!")
+    raise Exception("constructor error didn't occur")
+
+
 def test_model_forward():
-    model = YOLO()
-    model.new("yolov8n.yaml")
+    model = YOLO.new("yolov8n.yaml")
     img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512)
     model.forward(img)
     model(img)
 
 
 def test_model_info():
-    model = YOLO()
-    model.new("yolov8n.yaml")
+    model = YOLO.new("yolov8n.yaml")
     model.info()
-    model.load("best.pt")
+    model = model.load("best.pt")
     model.info(verbose=True)
 
 
 def test_model_fuse():
-    model = YOLO()
-    model.new("yolov8n.yaml")
+    model = YOLO.new("yolov8n.yaml")
     model.fuse()
     model.load("best.pt")
     model.fuse()
 
 
 def test_visualize_preds():
-    model = YOLO()
-    model.load("best.pt")
+    model = YOLO.load("best.pt")
     model.predict(source="ultralytics/assets")
 
 
 def test_val():
-    model = YOLO()
-    model.load("best.pt")
+    model = YOLO.load("best.pt")
     model.val(data="coco128.yaml", imgsz=32)
 
 
 def test_model_resume():
-    model = YOLO()
-    model.new("yolov8n.yaml")
+    model = YOLO.new("yolov8n.yaml")
     model.train(epochs=1, imgsz=32, data="coco128.yaml")
     try:
         model.resume(task="detect")
@@ -50,10 +54,9 @@ def test_model_resume():
 
 
 def test_model_train_pretrained():
-    model = YOLO()
-    model.load("best.pt")
+    model = YOLO.load("best.pt")
     model.train(data="coco128.yaml", epochs=1, imgsz=32)
-    model.new("yolov8n.yaml")
+    model = model.new("yolov8n.yaml")
     model.train(data="coco128.yaml", epochs=1, imgsz=32)
     img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512)
     model(img)

ultralytics/yolo/engine/model.py

@@ -5,7 +5,7 @@ from ultralytics import yolo  # noqa required for python usage
 from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights
 # from ultralytics.yolo.data.utils import check_dataset, check_dataset_yaml
 from ultralytics.yolo.engine.trainer import DEFAULT_CONFIG
-from ultralytics.yolo.utils import LOGGER
+from ultralytics.yolo.utils import HELP_MSG, LOGGER
 from ultralytics.yolo.utils.checks import check_yaml
 from ultralytics.yolo.utils.configs import get_config
 from ultralytics.yolo.utils.files import yaml_load
@@ -28,12 +28,16 @@ class YOLO:
     """
     Python interface which emulates a model-like behaviour by wrapping trainers.
     """
+    __init_key = object()
 
-    def __init__(self, type="v8") -> None:
+    def __init__(self, init_key=None, type="v8") -> None:
         """
         Args:
             type (str): Type/version of models to use
         """
+        if init_key != YOLO.__init_key:
+            raise Exception(HELP_MSG)
+
         self.type = type
         self.ModelClass = None
         self.TrainerClass = None
@@ -44,8 +48,10 @@ class YOLO:
         self.task = None
         self.ckpt = None
         self.overrides = {}
+        self.init_disabled = False
 
-    def new(self, cfg: str):
+    @classmethod
+    def new(cls, cfg: str):
         """
         Initializes a new model and infers the task type from the model definitions
 
@@ -55,12 +61,15 @@ class YOLO:
         cfg = check_yaml(cfg)  # check YAML
         with open(cfg, encoding='ascii', errors='ignore') as f:
             cfg = yaml.safe_load(f)  # model dict
-        self.task = self._guess_task_from_head(cfg["head"][-1][-2])
-        self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._guess_ops_from_task(
-            self.task)
-        self.model = self.ModelClass(cfg)  # initialize
+        obj = cls(init_key=cls.__init_key)
+        obj.task = obj._guess_task_from_head(cfg["head"][-1][-2])
+        obj.ModelClass, obj.TrainerClass, obj.ValidatorClass, obj.PredictorClass = obj._guess_ops_from_task(obj.task)
+        obj.model = obj.ModelClass(cfg)  # initialize
+
+        return obj
 
-    def load(self, weights: str):
+    @classmethod
+    def load(cls, weights: str):
         """
         Initializes a new model and infers the task type from the model head
 
@@ -68,15 +77,18 @@ class YOLO:
             weights (str): model checkpoint to be loaded
 
         """
-        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
+        obj = cls(init_key=cls.__init_key)
+        obj.ckpt = torch.load(weights, map_location="cpu")
+        obj.task = obj.ckpt["train_args"]["task"]
+        obj.overrides = dict(obj.ckpt["train_args"])
+        obj.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)
+        obj.ModelClass, obj.TrainerClass, obj.ValidatorClass, obj.PredictorClass = obj._guess_ops_from_task(
+            task=obj.task)
+        obj.model = attempt_load_weights(weights)
+
+        return obj
 
     def reset(self):
         """

ultralytics/yolo/utils/__init__.py

@@ -21,6 +21,32 @@ FONT = 'Arial.ttf'  # https://ultralytics.com/assets/Arial.ttf
 VERBOSE = str(os.getenv('YOLOv5_VERBOSE', True)).lower() == 'true'  # global verbose mode
 TQDM_BAR_FORMAT = '{l_bar}{bar:10}{r_bar}'  # tqdm bar format
 LOGGING_NAME = 'yolov5'
+HELP_MSG = \
+    """
+    Please refer to below Usage examples for help running YOLOv8
+    For help visit Ultralytics Community at https://community.ultralytics.com/
+    Submit bug reports to https//github.com/ultralytics/ultralytics
+
+    Install:
+        pip install ultralytics
+
+    Python usage:
+        from ultralytics import YOLO
+
+        model = YOLO.new('yolov8n.yaml')  # create a new model from scratch
+        model = YOLO.load('yolov8n.pt')  # load a pretrained model (recommended for best training results)
+        results = model.train(data='coco128.yaml')
+        results = model.val()
+        results = model.predict(source='bus.jpg')
+        success = model.export(format='onnx')
+
+    CLI usage:
+        yolo task=detect    mode=train     model=yolov8n.yaml ...
+                  classify       predict         yolov8n-cls.yaml
+                  segment        val             yolov8n-seg.yaml
+
+    For all arguments see https://github.com/ultralytics/ultralytics/blob/main/ultralytics/yolo/utils/configs/default.yaml
+    """
 
 # Settings
 # torch.set_printoptions(linewidth=320, precision=5, profile='long')