mirror of
https://github.com/THU-MIG/yolov10.git
synced 2025-05-24 05:55:51 +08:00
Simplify usage, remove YOLO.new() and YOLO.load()
This commit is contained in:
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
|
def _load(self, weights: str):
|
||||||
|
|
||||||
@classmethod
|
|
||||||
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…
x
Reference in New Issue
Block a user