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 <ericpedley@gmail.com>

README.md

@@ -198,16 +198,16 @@ See [Classification Docs](https://docs.ultralytics.com/tasks/classification/) fo
 
 <div align="center">
   <a href="https://roboflow.com/?ref=ultralytics">
-    <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-roboflow.png" width="10%" /></a>
+    <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-roboflow.png" width="10%" /></a>
   <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
   <a href="https://cutt.ly/yolov5-readme-clearml">
-    <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-clearml.png" width="10%" /></a>
+    <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-clearml.png" width="10%" /></a>
   <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
   <a href="https://bit.ly/yolov5-readme-comet">
-    <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-comet.png" width="10%" /></a>
+    <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-comet.png" width="10%" /></a>
   <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
   <a href="https://bit.ly/yolov5-neuralmagic">
-    <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-neuralmagic.png" width="10%" /></a>
+    <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-neuralmagic.png" width="10%" /></a>
 </div>
 
 |                                                           Roboflow                                                           |                                                            ClearML ⭐ NEW                                                            |                                                                        Comet ⭐ NEW                                                                         |                                           Neural Magic ⭐ NEW                                           |
@@ -232,7 +232,8 @@ on your experience. Thank you 🙏 to all our contributors!
 
 <!-- SVG image from https://opencollective.com/ultralytics/contributors.svg?width=990 -->
 
-<a href="https://github.com/ultralytics/ultralytics/graphs/contributors"><img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/image-contributors-1280.png"/></a>
+<a href="https://github.com/ultralytics/yolov5/graphs/contributors">
+<img src="https://github.com/ultralytics/assets/raw/main/im/image-contributors.png" /></a>
 
 ## <div align="center">License</div>
 

README.zh-CN.md

@@ -174,16 +174,16 @@ success = model.export(format="onnx")  # 将模型导出为 ONNX 格式
 
 <div align="center">
   <a href="https://roboflow.com/?ref=ultralytics">
-    <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-roboflow.png" width="10%" /></a>
+    <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-roboflow.png" width="10%" /></a>
   <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
   <a href="https://cutt.ly/yolov5-readme-clearml">
-    <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-clearml.png" width="10%" /></a>
+    <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-clearml.png" width="10%" /></a>
   <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
   <a href="https://bit.ly/yolov5-readme-comet">
-    <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-comet.png" width="10%" /></a>
+    <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-comet.png" width="10%" /></a>
   <img src="https://github.com/ultralytics/assets/raw/main/social/logo-transparent.png" width="15%" height="0" alt="" />
   <a href="https://bit.ly/yolov5-neuralmagic">
-    <img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/logo-neuralmagic.png" width="10%" /></a>
+    <img src="https://github.com/ultralytics/assets/raw/main/partners/logo-neuralmagic.png" width="10%" /></a>
 </div>
 
 |                                      Roboflow                                      |                                 ClearML ⭐ 新                                 |                                     Comet ⭐ 新                                     |                                    Neural Magic ⭐ 新                                    |
@@ -203,7 +203,8 @@ success = model.export(format="onnx")  # 将模型导出为 ONNX 格式
 
 <!-- SVG image from https://opencollective.com/ultralytics/contributors.svg?width=990 -->
 
-<a href="https://github.com/ultralytics/ultralytics/graphs/contributors"><img src="https://github.com/ultralytics/yolov5/releases/download/v1.0/image-contributors-1280.png"/></a>
+<a href="https://github.com/ultralytics/yolov5/graphs/contributors">
+<img src="https://github.com/ultralytics/assets/raw/main/im/image-contributors.png" /></a>
 
 ## <div align="center">License</div>
 

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",
+        "Operating System :: Microsoft :: Windows",],
-        "Development Status :: 4 - Beta",],
     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']})

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

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

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:

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,
+from ultralytics.yolo.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_PATH, LOGGER, ROOT, USER_CONFIG_DIR,
-                                    USER_CONFIG_DIR, IterableSimpleNamespace, __version__, colorstr, emojis, yaml_load,
+                                    IterableSimpleNamespace, __version__, checks, colorstr, yaml_load, yaml_print)
-                                    yaml_print)
-from ultralytics.yolo.utils.checks import check_yolo
 
 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")
 
 

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

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:
+        try:
-            {'.pt': self._load, '.yaml': self._new}[suffix](model)
+            if suffix == '.yaml':
-        else:
+                self._new(model)
-            raise NotImplementedError(f"'{suffix}' models not supported. Try a *.pt and *.yaml model, "
+            else:
-                                      "i.e. model='yolov8n.pt' or model='yolov8n.yaml'")
+                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
+        self.cfg = check_yaml(cfg)  # check YAML
-        cfg_dict = yaml_load(cfg, append_filename=True)  # model dict
+        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.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task()
-            self._assign_ops_from_task(self.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.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_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)
 
     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):
+    def _assign_ops_from_task(self):
-        model_class, train_lit, val_lit, pred_lit = MODEL_MAP[task]
+        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}"))

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()

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)

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
 

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)

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):

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)
+                                    nc=self.nc)
-        return p, 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
+        return p, proto
 
     def update_metrics(self, preds, batch):
         # Metrics