Simplify usage, remove YOLO.new() and YOLO.load()

single_channel
Glenn Jocher 2 years ago
parent 7b21a87da3
commit d76d7af566

@ -4,7 +4,7 @@ from ultralytics.yolo.utils import ROOT
if __name__ == "__main__": if __name__ == "__main__":
for m in list((ROOT / 'yolo/v8/models').rglob('*.yaml')): for m in list((ROOT / 'yolo/v8/models').rglob('*.yaml')):
try: try:
YOLO.new(m.name, verbose=True) YOLO(m.name, verbose=True)
except Exception as e: except Exception as e:
print(f'ERROR for {m}: {e}') print(f'ERROR for {m}: {e}')

@ -4,7 +4,7 @@ from ultralytics import YOLO
def test_model_init(): def test_model_init():
model = YOLO.new("yolov8n.yaml") model = YOLO("yolov8n.yaml")
model.info() model.info()
try: try:
YOLO() YOLO()
@ -14,38 +14,38 @@ def test_model_init():
def test_model_forward(): def test_model_forward():
model = YOLO.new("yolov8n.yaml") model = YOLO("yolov8n.yaml")
img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512) img = torch.rand(512 * 512 * 3).view(1, 3, 512, 512)
model.forward(img) model.forward(img)
model(img) model(img)
def test_model_info(): def test_model_info():
model = YOLO.new("yolov8n.yaml") model = YOLO("yolov8n.yaml")
model.info() model.info()
model = model.load("best.pt") model = model.load("best.pt")
model.info(verbose=True) model.info(verbose=True)
def test_model_fuse(): def test_model_fuse():
model = YOLO.new("yolov8n.yaml") model = YOLO("yolov8n.yaml")
model.fuse() model.fuse()
model.load("best.pt") model.load("best.pt")
model.fuse() model.fuse()
def test_visualize_preds(): def test_visualize_preds():
model = YOLO.load("best.pt") model = YOLO("best.pt")
model.predict(source="ultralytics/assets") model.predict(source="ultralytics/assets")
def test_val(): def test_val():
model = YOLO.load("best.pt") model = YOLO("best.pt")
model.val(data="coco128.yaml", imgsz=32) model.val(data="coco128.yaml", imgsz=32)
def test_model_resume(): def test_model_resume():
model = YOLO.new("yolov8n.yaml") model = YOLO("yolov8n.yaml")
model.train(epochs=1, imgsz=32, data="coco128.yaml") model.train(epochs=1, imgsz=32, data="coco128.yaml")
try: try:
model.resume(task="detect") model.resume(task="detect")
@ -54,7 +54,7 @@ def test_model_resume():
def test_model_train_pretrained(): def test_model_train_pretrained():
model = YOLO.load("best.pt") model = YOLO("best.pt")
model.train(data="coco128.yaml", epochs=1, imgsz=32) model.train(data="coco128.yaml", epochs=1, imgsz=32)
model = model.new("yolov8n.yaml") model = model.new("yolov8n.yaml")
model.train(data="coco128.yaml", epochs=1, imgsz=32) model.train(data="coco128.yaml", epochs=1, imgsz=32)
@ -83,7 +83,7 @@ def test_exports():
print(export_formats()) print(export_formats())
model = YOLO.new("yolov8n.yaml") model = YOLO("yolov8n.yaml")
model.export(format='torchscript') model.export(format='torchscript')
model.export(format='onnx') model.export(format='onnx')
model.export(format='openvino') model.export(format='openvino')

@ -23,7 +23,7 @@ Requirements:
Python: Python:
from ultralytics import YOLO from ultralytics import YOLO
model = YOLO.new('yolov8n.yaml') model = YOLO('yolov8n.yaml')
results = model.export(format='onnx') results = model.export(format='onnx')
CLI: CLI:

@ -1,4 +1,5 @@
import torch import torch
from pathlib import Path
from ultralytics import yolo # noqa from ultralytics import yolo # noqa
from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_weights
@ -27,19 +28,15 @@ class YOLO:
A python interface which emulates a model-like behaviour by wrapping trainers. A python interface which emulates a model-like behaviour by wrapping trainers.
""" """
__init_key = object() # used to ensure proper initialization
def __init__(self, init_key=None, type="v8") -> None: def __init__(self, model='yolov8n.yaml', type="v8") -> None:
""" """
Initializes the YOLO object. Initializes the YOLO object.
Args: Args:
init_key (object): used to ensure proper initialization. Defaults to None. model (str, Path): model to load or create
type (str): Type/version of models to use. Defaults to "v8". type (str): Type/version of models to use. Defaults to "v8".
""" """
if init_key != YOLO.__init_key:
raise SyntaxError(HELP_MSG)
self.type = type self.type = type
self.ModelClass = None # model class self.ModelClass = None # model class
self.TrainerClass = None # trainer class self.TrainerClass = None # trainer class
@ -53,8 +50,10 @@ class YOLO:
self.overrides = {} # overrides for trainer object self.overrides = {} # overrides for trainer object
self.init_disabled = False # disable model initialization self.init_disabled = False # disable model initialization
@classmethod # Load or create new YOLO model
def new(cls, cfg: str, verbose=True): {'.pt': self._load, '.yaml': self._new}[Path(model).suffix](model)
def _new(self, cfg: str, verbose=True):
""" """
Initializes a new model and infers the task type from the model definitions. Initializes a new model and infers the task type from the model definitions.
@ -64,34 +63,26 @@ class YOLO:
""" """
cfg = check_yaml(cfg) # check YAML cfg = check_yaml(cfg) # check YAML
cfg_dict = yaml_load(cfg) # model dict cfg_dict = yaml_load(cfg) # model dict
obj = cls(init_key=cls.__init_key) self.task = guess_task_from_head(cfg_dict["head"][-1][-2])
obj.task = guess_task_from_head(cfg_dict["head"][-1][-2]) self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = \
obj.ModelClass, obj.TrainerClass, obj.ValidatorClass, obj.PredictorClass = obj._guess_ops_from_task(obj.task) self._guess_ops_from_task(self.task)
obj.model = obj.ModelClass(cfg_dict, verbose=verbose) # initialize self.model = self.ModelClass(cfg_dict, verbose=verbose) # initialize
obj.cfg = cfg self.cfg = cfg
return obj
@classmethod def _load(self, weights: str):
def load(cls, weights: str):
""" """
Initializes a new model and infers the task type from the model head Initializes a new model and infers the task type from the model head
Args: Args:
weights (str): model checkpoint to be loaded weights (str): model checkpoint to be loaded
""" """
obj = cls(init_key=cls.__init_key) self.ckpt = torch.load(weights, map_location="cpu")
obj.ckpt = torch.load(weights, map_location="cpu") self.task = self.ckpt["train_args"]["task"]
obj.task = obj.ckpt["train_args"]["task"] self.overrides = dict(self.ckpt["train_args"])
obj.overrides = dict(obj.ckpt["train_args"]) self.overrides["device"] = '' # reset device
obj.overrides["device"] = '' # reset device self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = \
LOGGER.info("Device has been reset to ''") self._guess_ops_from_task(self.task)
self.model = attempt_load_weights(weights, fuse=False)
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): def reset(self):
""" """

@ -36,8 +36,8 @@ HELP_MSG = \
from ultralytics import YOLO from ultralytics import YOLO
model = YOLO.new('yolov8n.yaml') # create a new model from scratch model = YOLO('yolov8n.yaml') # build a new model from scratch
model = YOLO.load('yolov8n.pt') # load a pretrained model (recommended for best training results) model = YOLO('yolov8n.pt') # load a pretrained model (recommended for best training results)
results = model.train(data='coco128.yaml') # train the model results = model.train(data='coco128.yaml') # train the model
results = model.val() # evaluate model performance on the validation set results = model.val() # evaluate model performance on the validation set
results = model.predict(source='bus.jpg') # predict on an image results = model.predict(source='bus.jpg') # predict on an image

Loading…
Cancel
Save