# Ultralytics YOLO 🚀, AGPL-3.0 license import sys from pathlib import Path from typing import Union from ultralytics import yolo # noqa from ultralytics.nn.tasks import (ClassificationModel, DetectionModel, PoseModel, SegmentationModel, attempt_load_one_weight, guess_model_task, nn, yaml_model_load) from ultralytics.yolo.cfg import get_cfg from ultralytics.yolo.engine.exporter import Exporter from ultralytics.yolo.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, ROOT, callbacks, is_git_dir, yaml_load) from ultralytics.yolo.utils.checks import check_file, check_imgsz, check_pip_update_available, 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 TASK_MAP = { 'classify': [ ClassificationModel, yolo.v8.classify.ClassificationTrainer, yolo.v8.classify.ClassificationValidator, yolo.v8.classify.ClassificationPredictor], 'detect': [ DetectionModel, yolo.v8.detect.DetectionTrainer, yolo.v8.detect.DetectionValidator, yolo.v8.detect.DetectionPredictor], 'segment': [ SegmentationModel, yolo.v8.segment.SegmentationTrainer, yolo.v8.segment.SegmentationValidator, yolo.v8.segment.SegmentationPredictor], 'pose': [PoseModel, yolo.v8.pose.PoseTrainer, yolo.v8.pose.PoseValidator, yolo.v8.pose.PosePredictor]} class YOLO: """ 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 (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: Union[str, Path] = 'yolov8n.pt', task=None) -> None: """ Initializes the YOLO model. Args: model (Union[str, Path], optional): Path or name of the model to load or create. Defaults to 'yolov8n.pt'. task (Any, optional): Task type for the YOLO model. Defaults to None. """ self.callbacks = callbacks.get_default_callbacks() self.predictor = None # reuse predictor self.model = None # model object self.trainer = None # trainer object self.task = None # task type self.ckpt = None # if loaded from *.pt self.cfg = None # if loaded from *.yaml self.ckpt_path = None self.overrides = {} # overrides for trainer object self.metrics = None # validation/training metrics self.session = None # HUB session model = str(model).strip() # strip spaces # Check if Ultralytics HUB model from https://hub.ultralytics.com if self.is_hub_model(model): from ultralytics.hub.session import HUBTrainingSession self.session = HUBTrainingSession(model) model = self.session.model_file # Load or create new YOLO model suffix = Path(model).suffix if not suffix and Path(model).stem in GITHUB_ASSET_STEMS: model, suffix = Path(model).with_suffix('.pt'), '.pt' # add suffix, i.e. yolov8n -> yolov8n.pt if suffix == '.yaml': self._new(model, task) else: self._load(model, task) 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__}") @staticmethod def is_hub_model(model): return any(( model.startswith('https://hub.ultra'), # i.e. https://hub.ultralytics.com/models/MODEL_ID [len(x) for x in model.split('_')] == [42, 20], # APIKEY_MODELID len(model) == 20 and not Path(model).exists() and all(x not in model for x in './\\'))) # MODELID def _new(self, cfg: str, task=None, verbose=True): """ Initializes a new model and infers the task type from the model definitions. Args: cfg (str): model configuration file task (str) or (None): model task verbose (bool): display model info on load """ cfg_dict = yaml_model_load(cfg) self.cfg = cfg self.task = task or guess_model_task(cfg_dict) self.model = TASK_MAP[self.task][0](cfg_dict, verbose=verbose and RANK == -1) # build model self.overrides['model'] = self.cfg # 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=None): """ Initializes a new model and infers the task type from the model head. Args: weights (str): model checkpoint to be loaded task (str) or (None): model task """ suffix = Path(weights).suffix 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.model.args) self.ckpt_path = self.model.pt_path else: weights = check_file(weights) self.model, self.ckpt = weights, None self.task = task or guess_model_task(weights) self.ckpt_path = weights self.overrides['model'] = weights self.overrides['task'] = self.task def _check_is_pytorch_model(self): """ Raises TypeError is model is not a PyTorch model """ pt_str = isinstance(self.model, (str, Path)) and Path(self.model).suffix == '.pt' pt_module = isinstance(self.model, nn.Module) if not (pt_module or pt_str): raise TypeError(f"model='{self.model}' must be a *.pt PyTorch model, but is a different type. " f'PyTorch models can be used to train, val, predict and export, i.e. ' f"'yolo export model=yolov8n.pt', but exported formats like ONNX, TensorRT etc. only " f"support 'predict' and 'val' modes, i.e. 'yolo predict model=yolov8n.onnx'.") @smart_inference_mode() def reset_weights(self): """ Resets the model modules parameters to randomly initialized values, losing all training information. """ self._check_is_pytorch_model() for m in self.model.modules(): if hasattr(m, 'reset_parameters'): m.reset_parameters() for p in self.model.parameters(): p.requires_grad = True return self @smart_inference_mode() def load(self, weights='yolov8n.pt'): """ Transfers parameters with matching names and shapes from 'weights' to model. """ self._check_is_pytorch_model() if isinstance(weights, (str, Path)): weights, self.ckpt = attempt_load_one_weight(weights) self.model.load(weights) return self def info(self, verbose=True): """ Logs model info. Args: verbose (bool): Controls verbosity. """ self._check_is_pytorch_model() self.model.info(verbose=verbose) def fuse(self): self._check_is_pytorch_model() self.model.fuse() @smart_inference_mode() def predict(self, source=None, stream=False, **kwargs): """ Perform prediction using the YOLO model. Args: source (str | int | PIL | np.ndarray): The source of the image to make predictions on. Accepts all source types accepted by the YOLO model. stream (bool): Whether to stream the predictions or not. Defaults to False. **kwargs : Additional keyword arguments passed to the predictor. Check the 'configuration' section in the documentation for all available options. Returns: (List[ultralytics.yolo.engine.results.Results]): The prediction results. """ if source is None: source = ROOT / 'assets' if is_git_dir() else 'https://ultralytics.com/images/bus.jpg' LOGGER.warning(f"WARNING ⚠️ 'source' is missing. Using 'source={source}'.") is_cli = (sys.argv[0].endswith('yolo') or sys.argv[0].endswith('ultralytics')) and any( x in sys.argv for x in ('predict', 'track', 'mode=predict', 'mode=track')) overrides = self.overrides.copy() overrides['conf'] = 0.25 overrides.update(kwargs) # prefer kwargs overrides['mode'] = kwargs.get('mode', 'predict') assert overrides['mode'] in ['track', 'predict'] if not is_cli: overrides['save'] = kwargs.get('save', False) # do not save by default if called in Python if not self.predictor: self.task = overrides.get('task') or self.task self.predictor = TASK_MAP[self.task][3](overrides=overrides, _callbacks=self.callbacks) self.predictor.setup_model(model=self.model, verbose=is_cli) else: # only update args if predictor is already setup self.predictor.args = get_cfg(self.predictor.args, overrides) return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream) def track(self, source=None, stream=False, persist=False, **kwargs): """ Perform object tracking on the input source using the registered trackers. Args: source (str, optional): The input source for object tracking. Can be a file path or a video stream. stream (bool, optional): Whether the input source is a video stream. Defaults to False. persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False. **kwargs (optional): Additional keyword arguments for the tracking process. Returns: (List[ultralytics.yolo.engine.results.Results]): The tracking results. """ if not hasattr(self.predictor, 'trackers'): from ultralytics.tracker import register_tracker register_tracker(self, persist) # ByteTrack-based method needs low confidence predictions as input conf = kwargs.get('conf') or 0.1 kwargs['conf'] = conf kwargs['mode'] = 'track' return self.predict(source=source, stream=stream, **kwargs) @smart_inference_mode() def val(self, data=None, **kwargs): """ Validate a model on a given dataset . Args: data (str): The dataset to validate on. Accepts all formats accepted by yolo **kwargs : Any other args accepted by the validators. To see all args check 'configuration' section in docs """ overrides = self.overrides.copy() overrides['rect'] = True # rect batches as default overrides.update(kwargs) 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 = TASK_MAP[self.task][2](args=args, _callbacks=self.callbacks) validator(model=self.model) self.metrics = validator.metrics return validator.metrics @smart_inference_mode() def benchmark(self, **kwargs): """ Benchmark a model on all export formats. Args: **kwargs : Any other args accepted by the validators. To see all args check 'configuration' section in docs """ self._check_is_pytorch_model() from ultralytics.yolo.utils.benchmarks import benchmark overrides = self.model.args.copy() overrides.update(kwargs) overrides['mode'] = 'benchmark' overrides = {**DEFAULT_CFG_DICT, **overrides} # fill in missing overrides keys with defaults return benchmark(model=self, imgsz=overrides['imgsz'], half=overrides['half'], device=overrides['device']) def export(self, **kwargs): """ Export model. Args: **kwargs : Any other args accepted by the predictors. To see all args check 'configuration' section in docs """ self._check_is_pytorch_model() overrides = self.overrides.copy() overrides.update(kwargs) overrides['mode'] = 'export' args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides) args.task = self.task if args.imgsz == DEFAULT_CFG.imgsz: 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 return Exporter(overrides=args, _callbacks=self.callbacks)(model=self.model) def train(self, **kwargs): """ Trains the model on a given dataset. Args: **kwargs (Any): Any number of arguments representing the training configuration. """ self._check_is_pytorch_model() if self.session: # Ultralytics HUB session if any(kwargs): LOGGER.warning('WARNING ⚠️ using HUB training arguments, ignoring local training arguments.') kwargs = self.session.train_args self.session.check_disk_space() check_pip_update_available() overrides = self.overrides.copy() 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'])) 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.task = overrides.get('task') or self.task self.trainer = TASK_MAP[self.task][1](overrides=overrides, _callbacks=self.callbacks) 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 self.trainer.hub_session = self.session # attach optional HUB session self.trainer.train() # update model and cfg after training if RANK in (-1, 0): self.model, _ = attempt_load_one_weight(str(self.trainer.best)) self.overrides = self.model.args self.metrics = getattr(self.trainer.validator, 'metrics', None) # TODO: no metrics returned by DDP def to(self, device): """ Sends the model to the given device. Args: device (str): device """ self._check_is_pytorch_model() self.model.to(device) def tune(self, data: str, space: dict = None, grace_period: int = 10, gpu_per_trial: int = None, max_samples: int = 10, train_args: dict = {}): """ Runs hyperparameter tuning using Ray Tune. Args: data (str): The dataset to run the tuner on. space (dict, optional): The hyperparameter search space. Defaults to None. grace_period (int, optional): The grace period in epochs of the ASHA scheduler. Defaults to 10. gpu_per_trial (int, optional): The number of GPUs to allocate per trial. Defaults to None. max_samples (int, optional): The maximum number of trials to run. Defaults to 10. train_args (dict, optional): Additional arguments to pass to the `train()` method. Defaults to {}. Returns: A dictionary containing the results of the hyperparameter search. Raises: ModuleNotFoundError: If Ray Tune is not installed. """ try: from ultralytics.yolo.utils.tuner import (ASHAScheduler, RunConfig, WandbLoggerCallback, default_space, task_metric_map, tune) except ImportError: raise ModuleNotFoundError("Install Ray Tune: `pip install 'ray[tune]'`") try: import wandb from wandb import __version__ # noqa except ImportError: wandb = False def _tune(config): """ Trains the YOLO model with the specified hyperparameters and additional arguments. Args: config (dict): A dictionary of hyperparameters to use for training. Returns: None. """ self._reset_callbacks() config.update(train_args) self.train(**config) if not space: LOGGER.warning('WARNING: search space not provided. Using default search space') space = default_space space['data'] = data # Define the trainable function with allocated resources trainable_with_resources = tune.with_resources(_tune, {'cpu': 8, 'gpu': gpu_per_trial if gpu_per_trial else 0}) # Define the ASHA scheduler for hyperparameter search asha_scheduler = ASHAScheduler(time_attr='epoch', metric=task_metric_map[self.task], mode='max', max_t=train_args.get('epochs') or 100, grace_period=grace_period, reduction_factor=3) # Define the callbacks for the hyperparameter search tuner_callbacks = [WandbLoggerCallback(project='yolov8_tune') if wandb else None] # Create the Ray Tune hyperparameter search tuner tuner = tune.Tuner(trainable_with_resources, param_space=space, tune_config=tune.TuneConfig(scheduler=asha_scheduler, num_samples=max_samples), run_config=RunConfig(callbacks=tuner_callbacks, local_dir='./runs')) # Run the hyperparameter search tuner.fit() # Return the results of the hyperparameter search return tuner.get_results() @property def names(self): """ Returns class names of the loaded model. """ return self.model.names if hasattr(self.model, 'names') else None @property def device(self): """ Returns device if PyTorch model """ return next(self.model.parameters()).device if isinstance(self.model, nn.Module) else None @property def transforms(self): """ Returns transform of the loaded model. """ return self.model.transforms if hasattr(self.model, 'transforms') else None def add_callback(self, event: str, func): """ Add callback """ self.callbacks[event].append(func) @staticmethod def _reset_ckpt_args(args): 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} def _reset_callbacks(self): for event in callbacks.default_callbacks.keys(): self.callbacks[event] = [callbacks.default_callbacks[event][0]]