|
|
|
|
@ -9,22 +9,22 @@ from ultralytics.nn.tasks import (ClassificationModel, DetectionModel, Segmentat
|
|
|
|
|
guess_model_task, nn)
|
|
|
|
|
from ultralytics.yolo.cfg import get_cfg
|
|
|
|
|
from ultralytics.yolo.engine.exporter import Exporter
|
|
|
|
|
from ultralytics.yolo.utils import DEFAULT_CFG, DEFAULT_CFG_DICT, LOGGER, RANK, callbacks, yaml_load
|
|
|
|
|
from ultralytics.yolo.utils import DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, callbacks, yaml_load
|
|
|
|
|
from ultralytics.yolo.utils.checks import check_file, check_imgsz, check_yaml
|
|
|
|
|
from ultralytics.yolo.utils.downloads import GITHUB_ASSET_STEMS
|
|
|
|
|
from ultralytics.yolo.utils.torch_utils import smart_inference_mode
|
|
|
|
|
|
|
|
|
|
# Map head to model, trainer, validator, and predictor classes
|
|
|
|
|
MODEL_MAP = {
|
|
|
|
|
TASK_MAP = {
|
|
|
|
|
'classify': [
|
|
|
|
|
ClassificationModel, 'yolo.TYPE.classify.ClassificationTrainer', 'yolo.TYPE.classify.ClassificationValidator',
|
|
|
|
|
'yolo.TYPE.classify.ClassificationPredictor'],
|
|
|
|
|
ClassificationModel, yolo.v8.classify.ClassificationTrainer, yolo.v8.classify.ClassificationValidator,
|
|
|
|
|
yolo.v8.classify.ClassificationPredictor],
|
|
|
|
|
'detect': [
|
|
|
|
|
DetectionModel, 'yolo.TYPE.detect.DetectionTrainer', 'yolo.TYPE.detect.DetectionValidator',
|
|
|
|
|
'yolo.TYPE.detect.DetectionPredictor'],
|
|
|
|
|
DetectionModel, yolo.v8.detect.DetectionTrainer, yolo.v8.detect.DetectionValidator,
|
|
|
|
|
yolo.v8.detect.DetectionPredictor],
|
|
|
|
|
'segment': [
|
|
|
|
|
SegmentationModel, 'yolo.TYPE.segment.SegmentationTrainer', 'yolo.TYPE.segment.SegmentationValidator',
|
|
|
|
|
'yolo.TYPE.segment.SegmentationPredictor']}
|
|
|
|
|
SegmentationModel, yolo.v8.segment.SegmentationTrainer, yolo.v8.segment.SegmentationValidator,
|
|
|
|
|
yolo.v8.segment.SegmentationPredictor]}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class YOLO:
|
|
|
|
|
@ -33,52 +33,48 @@ class YOLO:
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
model (str, Path): Path to the model file to load or create.
|
|
|
|
|
type (str): Type/version of models to use. Defaults to "v8".
|
|
|
|
|
|
|
|
|
|
Attributes:
|
|
|
|
|
type (str): Type/version of models being used.
|
|
|
|
|
ModelClass (Any): Model class.
|
|
|
|
|
TrainerClass (Any): Trainer class.
|
|
|
|
|
ValidatorClass (Any): Validator class.
|
|
|
|
|
PredictorClass (Any): Predictor class.
|
|
|
|
|
predictor (Any): Predictor object.
|
|
|
|
|
model (Any): Model object.
|
|
|
|
|
trainer (Any): Trainer object.
|
|
|
|
|
task (str): Type of model task.
|
|
|
|
|
ckpt (Any): Checkpoint object if model loaded from *.pt file.
|
|
|
|
|
cfg (str): Model configuration if loaded from *.yaml file.
|
|
|
|
|
ckpt_path (str): Checkpoint file path.
|
|
|
|
|
overrides (dict): Overrides for trainer object.
|
|
|
|
|
metrics_data (Any): Data for metrics.
|
|
|
|
|
predictor (Any): The predictor object.
|
|
|
|
|
model (Any): The model object.
|
|
|
|
|
trainer (Any): The trainer object.
|
|
|
|
|
task (str): The type of model task.
|
|
|
|
|
ckpt (Any): The checkpoint object if the model loaded from *.pt file.
|
|
|
|
|
cfg (str): The model configuration if loaded from *.yaml file.
|
|
|
|
|
ckpt_path (str): The checkpoint file path.
|
|
|
|
|
overrides (dict): Overrides for the trainer object.
|
|
|
|
|
metrics_data (Any): The data for metrics.
|
|
|
|
|
|
|
|
|
|
Methods:
|
|
|
|
|
__call__(): Alias for predict method.
|
|
|
|
|
_new(cfg, verbose=True): Initializes a new model and infers the task type from the model definitions.
|
|
|
|
|
_load(weights): Initializes a new model and infers the task type from the model head.
|
|
|
|
|
_check_is_pytorch_model(): Raises TypeError if model is not a PyTorch model.
|
|
|
|
|
reset(): Resets the model modules.
|
|
|
|
|
info(verbose=False): Logs model info.
|
|
|
|
|
fuse(): Fuse model for faster inference.
|
|
|
|
|
predict(source=None, stream=False, **kwargs): Perform prediction using the YOLO model.
|
|
|
|
|
__call__(source=None, stream=False, **kwargs):
|
|
|
|
|
Alias for the predict method.
|
|
|
|
|
_new(cfg:str, verbose:bool=True) -> None:
|
|
|
|
|
Initializes a new model and infers the task type from the model definitions.
|
|
|
|
|
_load(weights:str, task:str='') -> None:
|
|
|
|
|
Initializes a new model and infers the task type from the model head.
|
|
|
|
|
_check_is_pytorch_model() -> None:
|
|
|
|
|
Raises TypeError if the model is not a PyTorch model.
|
|
|
|
|
reset() -> None:
|
|
|
|
|
Resets the model modules.
|
|
|
|
|
info(verbose:bool=False) -> None:
|
|
|
|
|
Logs the model info.
|
|
|
|
|
fuse() -> None:
|
|
|
|
|
Fuses the model for faster inference.
|
|
|
|
|
predict(source=None, stream=False, **kwargs) -> List[ultralytics.yolo.engine.results.Results]:
|
|
|
|
|
Performs prediction using the YOLO model.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
list(ultralytics.yolo.engine.results.Results): The prediction results.
|
|
|
|
|
list[ultralytics.yolo.engine.results.Results]: The prediction results.
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
def __init__(self, model='yolov8n.pt', type='v8') -> None:
|
|
|
|
|
def __init__(self, model='yolov8n.pt') -> None:
|
|
|
|
|
"""
|
|
|
|
|
Initializes the YOLO model.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
model (str, Path): model to load or create
|
|
|
|
|
type (str): Type/version of models to use. Defaults to "v8".
|
|
|
|
|
"""
|
|
|
|
|
self._reset_callbacks()
|
|
|
|
|
self.type = type
|
|
|
|
|
self.ModelClass = None # model class
|
|
|
|
|
self.TrainerClass = None # trainer class
|
|
|
|
|
self.ValidatorClass = None # validator class
|
|
|
|
|
self.PredictorClass = None # predictor class
|
|
|
|
|
self.predictor = None # reuse predictor
|
|
|
|
|
self.model = None # model object
|
|
|
|
|
self.trainer = None # trainer object
|
|
|
|
|
@ -101,6 +97,10 @@ class YOLO:
|
|
|
|
|
def __call__(self, source=None, stream=False, **kwargs):
|
|
|
|
|
return self.predict(source, stream, **kwargs)
|
|
|
|
|
|
|
|
|
|
def __getattr__(self, attr):
|
|
|
|
|
name = self.__class__.__name__
|
|
|
|
|
raise AttributeError(f"'{name}' object has no attribute '{attr}'. See valid attributes below.\n{self.__doc__}")
|
|
|
|
|
|
|
|
|
|
def _new(self, cfg: str, verbose=True):
|
|
|
|
|
"""
|
|
|
|
|
Initializes a new model and infers the task type from the model definitions.
|
|
|
|
|
@ -112,11 +112,15 @@ class YOLO:
|
|
|
|
|
self.cfg = check_yaml(cfg) # check YAML
|
|
|
|
|
cfg_dict = yaml_load(self.cfg, append_filename=True) # model dict
|
|
|
|
|
self.task = guess_model_task(cfg_dict)
|
|
|
|
|
self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task()
|
|
|
|
|
self.model = self.ModelClass(cfg_dict, verbose=verbose and RANK == -1) # initialize
|
|
|
|
|
self.model = TASK_MAP[self.task][0](cfg_dict, verbose=verbose and RANK == -1) # build model
|
|
|
|
|
self.overrides['model'] = self.cfg
|
|
|
|
|
|
|
|
|
|
def _load(self, weights: str):
|
|
|
|
|
# Below added to allow export from yamls
|
|
|
|
|
args = {**DEFAULT_CFG_DICT, **self.overrides} # combine model and default args, preferring model args
|
|
|
|
|
self.model.args = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS} # attach args to model
|
|
|
|
|
self.model.task = self.task
|
|
|
|
|
|
|
|
|
|
def _load(self, weights: str, task=''):
|
|
|
|
|
"""
|
|
|
|
|
Initializes a new model and infers the task type from the model head.
|
|
|
|
|
|
|
|
|
|
@ -127,8 +131,7 @@ class YOLO:
|
|
|
|
|
if suffix == '.pt':
|
|
|
|
|
self.model, self.ckpt = attempt_load_one_weight(weights)
|
|
|
|
|
self.task = self.model.args['task']
|
|
|
|
|
self.overrides = self.model.args
|
|
|
|
|
self._reset_ckpt_args(self.overrides)
|
|
|
|
|
self.overrides = self.model.args = self._reset_ckpt_args(self.model.args)
|
|
|
|
|
self.ckpt_path = self.model.pt_path
|
|
|
|
|
else:
|
|
|
|
|
weights = check_file(weights)
|
|
|
|
|
@ -136,7 +139,6 @@ class YOLO:
|
|
|
|
|
self.task = guess_model_task(weights)
|
|
|
|
|
self.ckpt_path = weights
|
|
|
|
|
self.overrides['model'] = weights
|
|
|
|
|
self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task()
|
|
|
|
|
|
|
|
|
|
def _check_is_pytorch_model(self):
|
|
|
|
|
"""
|
|
|
|
|
@ -189,12 +191,13 @@ class YOLO:
|
|
|
|
|
"""
|
|
|
|
|
overrides = self.overrides.copy()
|
|
|
|
|
overrides['conf'] = 0.25
|
|
|
|
|
overrides.update(kwargs)
|
|
|
|
|
overrides.update(kwargs) # prefer kwargs
|
|
|
|
|
overrides['mode'] = kwargs.get('mode', 'predict')
|
|
|
|
|
assert overrides['mode'] in ['track', 'predict']
|
|
|
|
|
overrides['save'] = kwargs.get('save', False) # not save files by default
|
|
|
|
|
if not self.predictor:
|
|
|
|
|
self.predictor = self.PredictorClass(overrides=overrides)
|
|
|
|
|
self.task = overrides.get('task') or self.task
|
|
|
|
|
self.predictor = TASK_MAP[self.task][3](overrides=overrides)
|
|
|
|
|
self.predictor.setup_model(model=self.model)
|
|
|
|
|
else: # only update args if predictor is already setup
|
|
|
|
|
self.predictor.args = get_cfg(self.predictor.args, overrides)
|
|
|
|
|
@ -226,12 +229,15 @@ class YOLO:
|
|
|
|
|
overrides['mode'] = 'val'
|
|
|
|
|
args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides)
|
|
|
|
|
args.data = data or args.data
|
|
|
|
|
if 'task' in overrides:
|
|
|
|
|
self.task = args.task
|
|
|
|
|
else:
|
|
|
|
|
args.task = self.task
|
|
|
|
|
if args.imgsz == DEFAULT_CFG.imgsz and not isinstance(self.model, (str, Path)):
|
|
|
|
|
args.imgsz = self.model.args['imgsz'] # use trained imgsz unless custom value is passed
|
|
|
|
|
args.imgsz = check_imgsz(args.imgsz, max_dim=1)
|
|
|
|
|
|
|
|
|
|
validator = self.ValidatorClass(args=args)
|
|
|
|
|
validator = TASK_MAP[self.task][2](args=args)
|
|
|
|
|
validator(model=self.model)
|
|
|
|
|
self.metrics_data = validator.metrics
|
|
|
|
|
|
|
|
|
|
@ -267,8 +273,7 @@ class YOLO:
|
|
|
|
|
args.imgsz = self.model.args['imgsz'] # use trained imgsz unless custom value is passed
|
|
|
|
|
if args.batch == DEFAULT_CFG.batch:
|
|
|
|
|
args.batch = 1 # default to 1 if not modified
|
|
|
|
|
exporter = Exporter(overrides=args)
|
|
|
|
|
return exporter(model=self.model)
|
|
|
|
|
return Exporter(overrides=args)(model=self.model)
|
|
|
|
|
|
|
|
|
|
def train(self, **kwargs):
|
|
|
|
|
"""
|
|
|
|
|
@ -282,15 +287,15 @@ class YOLO:
|
|
|
|
|
overrides.update(kwargs)
|
|
|
|
|
if kwargs.get('cfg'):
|
|
|
|
|
LOGGER.info(f"cfg file passed. Overriding default params with {kwargs['cfg']}.")
|
|
|
|
|
overrides = yaml_load(check_yaml(kwargs['cfg']), append_filename=True)
|
|
|
|
|
overrides['task'] = self.task
|
|
|
|
|
overrides = yaml_load(check_yaml(kwargs['cfg']))
|
|
|
|
|
overrides['mode'] = 'train'
|
|
|
|
|
if not overrides.get('data'):
|
|
|
|
|
raise AttributeError("Dataset required but missing, i.e. pass 'data=coco128.yaml'")
|
|
|
|
|
if overrides.get('resume'):
|
|
|
|
|
overrides['resume'] = self.ckpt_path
|
|
|
|
|
|
|
|
|
|
self.trainer = self.TrainerClass(overrides=overrides)
|
|
|
|
|
self.task = overrides.get('task') or self.task
|
|
|
|
|
self.trainer = TASK_MAP[self.task][1](overrides=overrides)
|
|
|
|
|
if not overrides.get('resume'): # manually set model only if not resuming
|
|
|
|
|
self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)
|
|
|
|
|
self.model = self.trainer.model
|
|
|
|
|
@ -311,13 +316,6 @@ class YOLO:
|
|
|
|
|
self._check_is_pytorch_model()
|
|
|
|
|
self.model.to(device)
|
|
|
|
|
|
|
|
|
|
def _assign_ops_from_task(self):
|
|
|
|
|
model_class, train_lit, val_lit, pred_lit = MODEL_MAP[self.task]
|
|
|
|
|
trainer_class = eval(train_lit.replace('TYPE', f'{self.type}'))
|
|
|
|
|
validator_class = eval(val_lit.replace('TYPE', f'{self.type}'))
|
|
|
|
|
predictor_class = eval(pred_lit.replace('TYPE', f'{self.type}'))
|
|
|
|
|
return model_class, trainer_class, validator_class, predictor_class
|
|
|
|
|
|
|
|
|
|
@property
|
|
|
|
|
def names(self):
|
|
|
|
|
"""
|
|
|
|
|
@ -357,9 +355,8 @@ class YOLO:
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def _reset_ckpt_args(args):
|
|
|
|
|
for arg in 'augment', 'verbose', 'project', 'name', 'exist_ok', 'resume', 'batch', 'epochs', 'cache', \
|
|
|
|
|
'save_json', 'half', 'v5loader', 'device', 'cfg', 'save', 'rect', 'plots', 'opset', 'simplify':
|
|
|
|
|
args.pop(arg, None)
|
|
|
|
|
include = {'imgsz', 'data', 'task', 'single_cls'} # only remember these arguments when loading a PyTorch model
|
|
|
|
|
return {k: v for k, v in args.items() if k in include}
|
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
|
|
def _reset_callbacks():
|
|
|
|
|
|
Glenn Jocher
2 years ago
committed by
GitHub