|
|
|
|
@ -9,76 +9,72 @@ 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: |
|
|
|
|
""" |
|
|
|
|
YOLO (You Only Look Once) object detection model. |
|
|
|
|
|
|
|
|
|
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. |
|
|
|
|
|
|
|
|
|
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. |
|
|
|
|
|
|
|
|
|
Returns: |
|
|
|
|
list(ultralytics.yolo.engine.results.Results): The prediction results. |
|
|
|
|
""" |
|
|
|
|
YOLO (You Only Look Once) object detection model. |
|
|
|
|
|
|
|
|
|
Args: |
|
|
|
|
model (str, Path): Path to the model file to load or create. |
|
|
|
|
|
|
|
|
|
Attributes: |
|
|
|
|
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__(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. |
|
|
|
|
""" |
|
|
|
|
|
|
|
|
|
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 |
|
|
|
|
args.task = self.task |
|
|
|
|
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