diff --git a/README.md b/README.md index 6e1e857..fea7133 100644 --- a/README.md +++ b/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. diff --git a/tests/test_model.py b/tests/test_model.py index 3146d81..91d143d 100644 --- a/tests/test_model.py +++ b/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) diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 15ffd12..ab59658 100644 --- a/ultralytics/yolo/engine/model.py +++ b/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): """ diff --git a/ultralytics/yolo/utils/__init__.py b/ultralytics/yolo/utils/__init__.py index 44fa103..76fc0c6 100644 --- a/ultralytics/yolo/utils/__init__.py +++ b/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')