From 977fd8f0b8764e83d7da964ba87477df6492735d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 11 Feb 2023 21:31:49 +0400 Subject: [PATCH] `ultralytics 8.0.35` TensorRT, ONNX and OpenVINO predict and val (#929) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Eric Pedley --- README.md | 11 +++--- README.zh-CN.md | 11 +++--- setup.py | 8 +++-- ultralytics/__init__.py | 2 +- ultralytics/nn/autobackend.py | 6 ++-- ultralytics/nn/tasks.py | 9 +++++ ultralytics/yolo/cfg/__init__.py | 32 ++++++++++-------- ultralytics/yolo/engine/exporter.py | 2 -- ultralytics/yolo/engine/model.py | 47 ++++++++++++++------------ ultralytics/yolo/engine/predictor.py | 4 +-- ultralytics/yolo/engine/validator.py | 2 +- ultralytics/yolo/utils/ops.py | 7 ++-- ultralytics/yolo/v8/detect/val.py | 3 +- ultralytics/yolo/v8/segment/predict.py | 4 +-- ultralytics/yolo/v8/segment/val.py | 9 +++-- 15 files changed, 88 insertions(+), 69 deletions(-) diff --git a/README.md b/README.md index ed6970e..3cdc743 100644 --- a/README.md +++ b/README.md @@ -198,16 +198,16 @@ See [Classification Docs](https://docs.ultralytics.com/tasks/classification/) fo
- + - + - + - +
| Roboflow | ClearML ⭐ NEW | Comet ⭐ NEW | Neural Magic ⭐ NEW | @@ -232,7 +232,8 @@ on your experience. Thank you 🙏 to all our contributors! - + + ##
License
diff --git a/README.zh-CN.md b/README.zh-CN.md index 5650a46..e2ee93f 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -174,16 +174,16 @@ success = model.export(format="onnx") # 将模型导出为 ONNX 格式
- + - + - + - +
| Roboflow | ClearML ⭐ 新 | Comet ⭐ 新 | Neural Magic ⭐ 新 | @@ -203,7 +203,8 @@ success = model.export(format="onnx") # 将模型导出为 ONNX 格式 - + + ##
License
diff --git a/setup.py b/setup.py index 4d3563c..35d85d9 100644 --- a/setup.py +++ b/setup.py @@ -22,7 +22,7 @@ def get_version(): setup( name="ultralytics", # name of pypi package version=get_version(), # version of pypi package - python_requires=">=3.7,<=3.11", + python_requires=">=3.7", license='GPL-3.0', description='Ultralytics YOLOv8', long_description=README, @@ -41,7 +41,9 @@ setup( 'dev': ['check-manifest', 'pytest', 'pytest-cov', 'coverage', 'mkdocs', 'mkdocstrings[python]', 'mkdocs-material']}, classifiers=[ + "Development Status :: 4 - Beta", "Intended Audience :: Developers", + "Intended Audience :: Education", "Intended Audience :: Science/Research", "License :: OSI Approved :: GNU General Public License v3 (GPLv3)", "Programming Language :: Python :: 3", @@ -49,14 +51,14 @@ setup( "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", "Topic :: Software Development", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Scientific/Engineering :: Image Recognition", "Operating System :: POSIX :: Linux", "Operating System :: MacOS", - "Operating System :: Microsoft :: Windows", - "Development Status :: 4 - Beta",], + "Operating System :: Microsoft :: Windows",], keywords="machine-learning, deep-learning, vision, ML, DL, AI, YOLO, YOLOv3, YOLOv5, YOLOv8, HUB, Ultralytics", entry_points={ 'console_scripts': ['yolo = ultralytics.yolo.cfg:entrypoint', 'ultralytics = ultralytics.yolo.cfg:entrypoint']}) diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index d9bdbec..6dfa993 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -1,6 +1,6 @@ # Ultralytics YOLO 🚀, GPL-3.0 license -__version__ = "8.0.34" +__version__ = "8.0.35" from ultralytics.yolo.engine.model import YOLO from ultralytics.yolo.utils.checks import check_yolo as checks diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py index 9248735..5a2bd4f 100644 --- a/ultralytics/nn/autobackend.py +++ b/ultralytics/nn/autobackend.py @@ -13,7 +13,7 @@ import torch.nn as nn from PIL import Image from ultralytics.yolo.utils import LOGGER, ROOT, yaml_load -from ultralytics.yolo.utils.checks import check_requirements, check_suffix, check_version +from ultralytics.yolo.utils.checks import check_requirements, check_suffix, check_version, check_yaml from ultralytics.yolo.utils.downloads import attempt_download_asset, is_url from ultralytics.yolo.utils.ops import xywh2xyxy @@ -38,7 +38,7 @@ class AutoBackend(nn.Module): weights (str): The path to the weights file. Default: 'yolov8n.pt' device (torch.device): The device to run the model on. dnn (bool): Use OpenCV's DNN module for inference if True, defaults to False. - data (dict): Additional data, optional + data (str), (Path): Additional data.yaml file for class names, optional fp16 (bool): If True, use half precision. Default: False fuse (bool): Whether to fuse the model or not. Default: True @@ -237,7 +237,7 @@ class AutoBackend(nn.Module): # class names if 'names' not in locals(): # names missing - names = yaml_load(data)['names'] if data else {i: f'class{i}' for i in range(999)} # assign default + names = yaml_load(check_yaml(data))['names'] if data else {i: f'class{i}' for i in range(999)} # assign names = check_class_names(names) self.__dict__.update(locals()) # assign all variables to self diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index 2cf7833..f0d80c4 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -2,6 +2,7 @@ import contextlib from copy import deepcopy +from pathlib import Path import thop import torch @@ -490,6 +491,14 @@ def guess_model_task(model): with contextlib.suppress(Exception): cfg = eval(x) break + elif isinstance(model, (str, Path)): + model = str(model) + if '-seg' in model: + return "segment" + elif '-cls' in model: + return "classify" + else: + return "detect" # Guess from YAML dictionary if cfg: diff --git a/ultralytics/yolo/cfg/__init__.py b/ultralytics/yolo/cfg/__init__.py index ec55e6d..404d4f7 100644 --- a/ultralytics/yolo/cfg/__init__.py +++ b/ultralytics/yolo/cfg/__init__.py @@ -8,10 +8,8 @@ from pathlib import Path from types import SimpleNamespace from typing import Dict, List, Union -from ultralytics.yolo.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_PATH, LOGGER, PREFIX, ROOT, - USER_CONFIG_DIR, IterableSimpleNamespace, __version__, colorstr, emojis, yaml_load, - yaml_print) -from ultralytics.yolo.utils.checks import check_yolo +from ultralytics.yolo.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_PATH, LOGGER, ROOT, USER_CONFIG_DIR, + IterableSimpleNamespace, __version__, checks, colorstr, yaml_load, yaml_print) CLI_HELP_MSG = \ """ @@ -83,7 +81,7 @@ def cfg2dict(cfg): return cfg -def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace] = DEFAULT_CFG, overrides: Dict = None): +def get_cfg(cfg: Union[str, Path, Dict, SimpleNamespace] = DEFAULT_CFG_DICT, overrides: Dict = None): """ Load and merge configuration data from a file or dictionary. @@ -198,17 +196,23 @@ def entrypoint(debug=''): LOGGER.info(CLI_HELP_MSG) return - # Add tasks, modes, special, and special with dash keys, i.e. -help, --help + # Define tasks and modes tasks = 'detect', 'segment', 'classify' modes = 'train', 'val', 'predict', 'export' + + # Define special commands special = { 'help': lambda: LOGGER.info(CLI_HELP_MSG), - 'checks': check_yolo, + 'checks': checks.check_yolo, 'version': lambda: LOGGER.info(__version__), 'settings': lambda: yaml_print(USER_CONFIG_DIR / 'settings.yaml'), 'cfg': lambda: yaml_print(DEFAULT_CFG_PATH), 'copy-cfg': copy_default_cfg} - FULL_ARGS_DICT = {**DEFAULT_CFG_DICT, **{k: None for k in tasks}, **{k: None for k in modes}, **special} + full_args_dict = {**DEFAULT_CFG_DICT, **{k: None for k in tasks}, **{k: None for k in modes}, **special} + + # Define common mis-uses of special commands, i.e. -h, -help, --help + special.update({k[0]: v for k, v in special.items()}) # singular + special.update({k[:-1]: v for k, v in special.items() if len(k) > 1 and k.endswith('s')}) # singular special = {**special, **{f'-{k}': v for k, v in special.items()}, **{f'--{k}': v for k, v in special.items()}} overrides = {} # basic overrides, i.e. imgsz=320 @@ -219,7 +223,7 @@ def entrypoint(debug=''): k, v = a.split('=', 1) # split on first '=' sign assert v, f"missing '{k}' value" if k == 'cfg': # custom.yaml passed - LOGGER.info(f"{PREFIX}Overriding {DEFAULT_CFG_PATH} with {v}") + LOGGER.info(f"Overriding {DEFAULT_CFG_PATH} with {v}") overrides = {k: val for k, val in yaml_load(v).items() if k != 'cfg'} else: if v.lower() == 'none': @@ -233,7 +237,7 @@ def entrypoint(debug=''): v = eval(v) overrides[k] = v except (NameError, SyntaxError, ValueError, AssertionError) as e: - check_cfg_mismatch(FULL_ARGS_DICT, {a: ""}, e) + check_cfg_mismatch(full_args_dict, {a: ""}, e) elif a in tasks: overrides['task'] = a @@ -248,7 +252,7 @@ def entrypoint(debug=''): raise SyntaxError(f"'{colorstr('red', 'bold', a)}' is a valid YOLO argument but is missing an '=' sign " f"to set its value, i.e. try '{a}={DEFAULT_CFG_DICT[a]}'\n{CLI_HELP_MSG}") else: - check_cfg_mismatch(FULL_ARGS_DICT, {a: ""}) + check_cfg_mismatch(full_args_dict, {a: ""}) # Defaults task2model = dict(detect='yolov8n.pt', segment='yolov8n-seg.pt', classify='yolov8n-cls.pt') @@ -261,9 +265,9 @@ def entrypoint(debug=''): LOGGER.warning(f"WARNING ⚠️ 'mode' is missing. Valid modes are {modes}. Using default 'mode={mode}'.") elif mode not in modes: if mode != 'checks': - raise ValueError(emojis(f"ERROR ❌ Invalid 'mode={mode}'. Valid modes are {modes}.")) + raise ValueError(f"Invalid 'mode={mode}'. Valid modes are {modes}.") LOGGER.warning("WARNING ⚠️ 'yolo mode=checks' is deprecated. Use 'yolo checks' instead.") - check_yolo() + checks.check_yolo() return # Model @@ -304,7 +308,7 @@ def entrypoint(debug=''): def copy_default_cfg(): new_file = Path.cwd() / DEFAULT_CFG_PATH.name.replace('.yaml', '_copy.yaml') shutil.copy2(DEFAULT_CFG_PATH, new_file) - LOGGER.info(f"{PREFIX}{DEFAULT_CFG_PATH} copied to {new_file}\n" + LOGGER.info(f"{DEFAULT_CFG_PATH} copied to {new_file}\n" f"Example YOLO command with this new custom cfg:\n yolo cfg='{new_file}' imgsz=320 batch=8") diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py index 053ba52..005a504 100644 --- a/ultralytics/yolo/engine/exporter.py +++ b/ultralytics/yolo/engine/exporter.py @@ -161,8 +161,6 @@ class Exporter: # Checks model.names = check_class_names(model.names) - # if self.args.batch == model.args['batch_size']: # user has not modified training batch_size - self.args.batch = 1 self.imgsz = check_imgsz(self.args.imgsz, stride=model.stride, min_dim=2) # check image size if model.task == 'classify': self.args.nms = self.args.agnostic_nms = False diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 7b0e61f..0bab311 100644 --- a/ultralytics/yolo/engine/model.py +++ b/ultralytics/yolo/engine/model.py @@ -57,15 +57,17 @@ class YOLO: self.overrides = {} # overrides for trainer object # Load or create new YOLO model - load_methods = {'.pt': self._load, '.yaml': self._new} 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 in load_methods: - {'.pt': self._load, '.yaml': self._new}[suffix](model) - else: - raise NotImplementedError(f"'{suffix}' models not supported. Try a *.pt and *.yaml model, " - "i.e. model='yolov8n.pt' or model='yolov8n.yaml'") + try: + if suffix == '.yaml': + self._new(model) + else: + self._load(model) + except Exception as e: + raise NotImplementedError(f"Unable to load model='{model}'. " + f"As an example try model='yolov8n.pt' or model='yolov8n.yaml'") from e def __call__(self, source=None, stream=False, **kwargs): return self.predict(source, stream, **kwargs) @@ -78,13 +80,11 @@ class YOLO: cfg (str): model configuration file verbose (bool): display model info on load """ - cfg = check_yaml(cfg) # check YAML - cfg_dict = yaml_load(cfg, append_filename=True) # model dict + 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.task) + self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task() self.model = self.ModelClass(cfg_dict, verbose=verbose) # initialize - self.cfg = cfg def _load(self, weights: str): """ @@ -93,13 +93,17 @@ class YOLO: Args: weights (str): model checkpoint to be loaded """ - self.model, self.ckpt = attempt_load_one_weight(weights) + 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.overrides) + else: + self.model, self.ckpt = weights, None + self.task = guess_model_task(weights) self.ckpt_path = weights - self.task = self.model.args["task"] - self.overrides = self.model.args - self._reset_ckpt_args(self.overrides) - self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = \ - self._assign_ops_from_task(self.task) + self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task() def reset(self): """ @@ -166,7 +170,7 @@ class YOLO: args = get_cfg(cfg=DEFAULT_CFG, overrides=overrides) args.data = data or args.data args.task = self.task - if args.imgsz == DEFAULT_CFG.imgsz: + 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) @@ -189,7 +193,8 @@ class YOLO: 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 exporter = Exporter(overrides=args) exporter(model=self.model) @@ -231,8 +236,8 @@ class YOLO: """ self.model.to(device) - def _assign_ops_from_task(self, task): - model_class, train_lit, val_lit, pred_lit = MODEL_MAP[task] + def _assign_ops_from_task(self): + model_class, train_lit, val_lit, pred_lit = MODEL_MAP[self.task] # warning: eval is unsafe. Use with caution trainer_class = eval(train_lit.replace("TYPE", f"{self.type}")) validator_class = eval(val_lit.replace("TYPE", f"{self.type}")) diff --git a/ultralytics/yolo/engine/predictor.py b/ultralytics/yolo/engine/predictor.py index 59dd85c..77382a7 100644 --- a/ultralytics/yolo/engine/predictor.py +++ b/ultralytics/yolo/engine/predictor.py @@ -146,7 +146,7 @@ class BasePredictor: (self.save_dir / 'labels' if self.args.save_txt else self.save_dir).mkdir(parents=True, exist_ok=True) # warmup model if not self.done_warmup: - self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else self.bs, 3, *self.imgsz)) + self.model.warmup(imgsz=(1 if self.model.pt or self.model.triton else self.dataset.bs, 3, *self.imgsz)) self.done_warmup = True self.seen, self.windows, self.dt, self.batch = 0, [], (ops.Profile(), ops.Profile(), ops.Profile()), None @@ -218,7 +218,7 @@ class BasePredictor: cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) # allow window resize (Linux) cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0]) cv2.imshow(str(p), im0) - cv2.waitKey(1) # 1 millisecond + cv2.waitKey(500 if self.batch[4].startswith('image') else 1) # 1 millisecond def save_preds(self, vid_cap, idx, save_path): im0 = self.annotator.result() diff --git a/ultralytics/yolo/engine/validator.py b/ultralytics/yolo/engine/validator.py index dfd9461..4bc5f3f 100644 --- a/ultralytics/yolo/engine/validator.py +++ b/ultralytics/yolo/engine/validator.py @@ -95,7 +95,7 @@ class BaseValidator: assert model is not None, "Either trainer or model is needed for validation" self.device = select_device(self.args.device, self.args.batch) self.args.half &= self.device.type != 'cpu' - model = AutoBackend(model, device=self.device, dnn=self.args.dnn, fp16=self.args.half) + model = AutoBackend(model, device=self.device, dnn=self.args.dnn, data=self.args.data, fp16=self.args.half) self.model = model stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine imgsz = check_imgsz(self.args.imgsz, stride=stride) diff --git a/ultralytics/yolo/utils/ops.py b/ultralytics/yolo/utils/ops.py index 4e213ec..5c95684 100644 --- a/ultralytics/yolo/utils/ops.py +++ b/ultralytics/yolo/utils/ops.py @@ -138,7 +138,7 @@ def non_max_suppression( multi_label=False, labels=(), max_det=300, - nm=0, # number of masks + nc=0, # number of classes (optional) ): """ Perform non-maximum suppression (NMS) on a set of boxes, with support for masks and multiple labels per box. @@ -159,7 +159,7 @@ def non_max_suppression( list contains the apriori labels for a given image. The list should be in the format output by a dataloader, with each label being a tuple of (class_index, x1, y1, x2, y2). max_det (int): The maximum number of boxes to keep after NMS. - nm (int): The number of masks output by the model. + nc (int): (optional) The number of classes output by the model. Any indices after this will be considered masks. Returns: (List[torch.Tensor]): A list of length batch_size, where each element is a tensor of @@ -178,7 +178,8 @@ def non_max_suppression( if mps: # MPS not fully supported yet, convert tensors to CPU before NMS prediction = prediction.cpu() bs = prediction.shape[0] # batch size - nc = prediction.shape[1] - nm - 4 # number of classes + nc = nc or (prediction.shape[1] - 4) # number of classes + nm = prediction.shape[1] - nc - 4 mi = 4 + nc # mask start index xc = prediction[:, 4:mi].amax(1) > conf_thres # candidates diff --git a/ultralytics/yolo/v8/detect/val.py b/ultralytics/yolo/v8/detect/val.py index b5127b3..bc6a306 100644 --- a/ultralytics/yolo/v8/detect/val.py +++ b/ultralytics/yolo/v8/detect/val.py @@ -40,13 +40,12 @@ class DetectionValidator(BaseValidator): return batch def init_metrics(self, model): - head = model.model[-1] if self.training else model.model.model[-1] val = self.data.get(self.args.split, '') # validation path self.is_coco = isinstance(val, str) and val.endswith(f'coco{os.sep}val2017.txt') # is COCO dataset self.class_map = ops.coco80_to_coco91_class() if self.is_coco else list(range(1000)) self.args.save_json |= self.is_coco and not self.training # run on final val if training COCO - self.nc = head.nc self.names = model.names + self.nc = len(model.names) self.metrics.names = self.names self.metrics.plot = self.args.plots self.confusion_matrix = ConfusionMatrix(nc=self.nc) diff --git a/ultralytics/yolo/v8/segment/predict.py b/ultralytics/yolo/v8/segment/predict.py index 4b5cd66..97ec3fb 100644 --- a/ultralytics/yolo/v8/segment/predict.py +++ b/ultralytics/yolo/v8/segment/predict.py @@ -17,10 +17,10 @@ class SegmentationPredictor(DetectionPredictor): self.args.iou, agnostic=self.args.agnostic_nms, max_det=self.args.max_det, - nm=32, + nc=len(self.model.names), classes=self.args.classes) results = [] - proto = preds[1][-1] + proto = preds[1][-1] if len(preds[1]) == 3 else preds[1] # second output is len 3 if pt, but only 1 if exported for i, pred in enumerate(p): shape = orig_img[i].shape if isinstance(orig_img, list) else orig_img.shape if not len(pred): diff --git a/ultralytics/yolo/v8/segment/val.py b/ultralytics/yolo/v8/segment/val.py index e0d35f8..556ac1f 100644 --- a/ultralytics/yolo/v8/segment/val.py +++ b/ultralytics/yolo/v8/segment/val.py @@ -28,14 +28,12 @@ class SegmentationValidator(DetectionValidator): return batch def init_metrics(self, model): - head = model.model[-1] if self.training else model.model.model[-1] val = self.data.get(self.args.split, '') # validation path self.is_coco = isinstance(val, str) and val.endswith(f'coco{os.sep}val2017.txt') # is COCO dataset self.class_map = ops.coco80_to_coco91_class() if self.is_coco else list(range(1000)) self.args.save_json |= self.is_coco and not self.training # run on final val if training COCO - self.nc = head.nc - self.nm = head.nm if hasattr(head, "nm") else 32 self.names = model.names + self.nc = len(model.names) self.metrics.names = self.names self.metrics.plot = self.args.plots self.confusion_matrix = ConfusionMatrix(nc=self.nc) @@ -61,8 +59,9 @@ class SegmentationValidator(DetectionValidator): multi_label=True, agnostic=self.args.single_cls, max_det=self.args.max_det, - nm=self.nm) - return p, preds[1][-1] + nc=self.nc) + proto = preds[1][-1] if len(preds[1]) == 3 else preds[1] # second output is len 3 if pt, but only 1 if exported + return p, proto def update_metrics(self, preds, batch): # Metrics