HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.0.18 new python callbacks and minor fixes (#580)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jeroen Rombouts <36196499+jarombouts@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ayush Chaurasia 2023-01-23 23:01:29 +05:30 • committed by GitHub
parent e9ab157330
commit 936414c615
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
24 changed files with 136 additions and 106 deletions
Download patch file Download diff file Expand all files Collapse all files

12
ultralytics/yolo/engine/model.py
View file

@ -6,9 +6,9 @@ from ultralytics import yolo # noqa
from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, attempt_load_one_weight
from ultralytics.yolo.cfg import get_cfg
from ultralytics.yolo.engine.exporter import Exporter
from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, yaml_load
from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, callbacks, yaml_load
from ultralytics.yolo.utils.checks import check_yaml
from ultralytics.yolo.utils.torch_utils import guess_task_from_head, smart_inference_mode
from ultralytics.yolo.utils.torch_utils import guess_task_from_model_yaml, smart_inference_mode
# Map head to model, trainer, validator, and predictor classes
MODEL_MAP = {
@ -68,7 +68,7 @@ class YOLO:
"""
cfg = check_yaml(cfg) # check YAML
cfg_dict = yaml_load(cfg, append_filename=True) # model dict
self.task = guess_task_from_head(cfg_dict["head"][-1][-2])
self.task = guess_task_from_model_yaml(cfg_dict)
self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = \
self._guess_ops_from_task(self.task)
self.model = self.ModelClass(cfg_dict, verbose=verbose) # initialize
@ -228,6 +228,12 @@ class YOLO:
"""
return self.model.names
def add_callback(self, event: str, func):
"""
Add callback
"""
callbacks.default_callbacks[event].append(func)
@staticmethod
def _reset_ckpt_args(args):
args.pop("project", None)