Update YOLOv5 YAMLs to 'u' YAMLs (#800)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher
2023-02-04 19:54:34 +04:00
committed by GitHub
parent 0d182e80f1
commit 21ae321bc2
17 changed files with 68 additions and 46 deletions

View File

@ -70,7 +70,7 @@ from ultralytics.nn.modules import Detect, Segment
from ultralytics.nn.tasks import ClassificationModel, DetectionModel, SegmentationModel, guess_model_task
from ultralytics.yolo.cfg import get_cfg
from ultralytics.yolo.data.dataloaders.stream_loaders import LoadImages
from ultralytics.yolo.data.utils import check_det_dataset
from ultralytics.yolo.data.utils import check_det_dataset, IMAGENET_MEAN, IMAGENET_STD
from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, callbacks, colorstr, get_default_args, yaml_save
from ultralytics.yolo.utils.checks import check_imgsz, check_requirements, check_version, check_yaml
from ultralytics.yolo.utils.files import file_size
@ -185,8 +185,8 @@ class Exporter:
if self.args.half and not coreml and not xml:
im, model = im.half(), model.half() # to FP16
shape = tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape
LOGGER.info(
f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with input shape {tuple(im.shape)} and "
f"output shape {shape} ({file_size(file):.1f} MB)")
# Warnings
warnings.filterwarnings('ignore', category=torch.jit.TracerWarning) # suppress TracerWarning
@ -384,12 +384,18 @@ class Exporter:
LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...')
f = self.file.with_suffix('.mlmodel')
task = self.model.task
if self.model.task == 'classify':
bias = [-x for x in IMAGENET_MEAN]
scale = 1 / 255 / (sum(IMAGENET_STD) / 3)
classifier_config = ct.ClassifierConfig(list(self.model.names.values()))
else:
bias = [0.0, 0.0, 0.0]
scale = 1 / 255
classifier_config = None
model = iOSModel(self.model, self.im).eval() if self.args.nms else self.model
ts = torch.jit.trace(model, self.im, strict=False) # TorchScript model
classifier_config = ct.ClassifierConfig(list(model.names.values())) if task == 'classify' else None
ct_model = ct.convert(ts,
inputs=[ct.ImageType('image', shape=self.im.shape, scale=1 / 255, bias=[0, 0, 0])],
inputs=[ct.ImageType('image', shape=self.im.shape, scale=scale, bias=bias)],
classifier_config=classifier_config)
bits, mode = (8, 'kmeans_lut') if self.args.int8 else (16, 'linear') if self.args.half else (32, None)
if bits < 32:

View File

@ -162,6 +162,8 @@ 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:
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)
@ -180,6 +182,8 @@ class YOLO:
overrides.update(kwargs)
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
exporter = Exporter(overrides=args)
exporter(model=self.model)

View File

@ -120,9 +120,6 @@ class BasePredictor:
pass
def setup_source(self, source):
if not self.model:
raise Exception("Model not initialized!")
self.imgsz = check_imgsz(self.args.imgsz, stride=self.model.stride, min_dim=2) # check image size
self.dataset = load_inference_source(source=source,
transforms=getattr(self.model.model, 'transforms', None),
@ -190,6 +187,10 @@ class BasePredictor:
if self.args.verbose:
LOGGER.info(f"{s}{'' if len(preds) else '(no detections), '}{self.dt[1].dt * 1E3:.1f}ms")
# Release assets
if isinstance(self.vid_writer[-1], cv2.VideoWriter):
self.vid_writer[-1].release() # release final video writer
# Print results
if self.args.verbose and self.seen:
t = tuple(x.t / self.seen * 1E3 for x in self.dt) # speeds per image

View File

@ -5,6 +5,7 @@ import inspect
import math
import os
import platform
import re
import shutil
import urllib
from pathlib import Path
@ -67,12 +68,13 @@ def check_imgsz(imgsz, stride=32, min_dim=1, max_dim=2, floor=0):
f"Valid imgsz types are int i.e. 'imgsz=640' or list i.e. 'imgsz=[640,640]'")
# Apply max_dim
if max_dim == 1:
LOGGER.warning(f"WARNING ⚠️ 'train' and 'val' imgsz types must be integer, updating to 'imgsz={max(imgsz)}'. "
f"'predict' and 'export' imgsz may be list or integer, "
f"i.e. 'yolo export imgsz=640,480' or 'yolo export imgsz=640'")
if len(imgsz) > max_dim:
msg = "'train' and 'val' imgsz must be an integer, while 'predict' and 'export' imgsz may be a [h, w] list " \
"or an integer, i.e. 'yolo export imgsz=640,480' or 'yolo export imgsz=640'"
if max_dim != 1:
raise ValueError(f"imgsz={imgsz} is not a valid image size. {msg}")
LOGGER.warning(f"WARNING ⚠️ updating to 'imgsz={max(imgsz)}'. {msg}")
imgsz = [max(imgsz)]
# Make image size a multiple of the stride
sz = [max(math.ceil(x / stride) * stride, floor) for x in imgsz]
@ -220,10 +222,24 @@ def check_suffix(file='yolov8n.pt', suffix=('.pt',), msg=''):
assert s in suffix, f"{msg}{f} acceptable suffix is {suffix}"
def check_yolov5u_filename(file: str):
# Replace legacy YOLOv5 filenames with updated YOLOv5u filenames
if 'yolov3' in file or 'yolov5' in file and 'u' not in file:
original_file = file
file = re.sub(r"(.*yolov5([nsmlx]))\.", "\\1u.", file) # i.e. yolov5n.pt -> yolov5nu.pt
file = re.sub(r"(.*yolov3(|-tiny|-spp))\.", "\\1u.", file) # i.e. yolov3-spp.pt -> yolov3-sppu.pt
if file != original_file:
LOGGER.info(f"PRO TIP 💡 Replace 'model={original_file}' with new 'model={file}'.\nYOLOv5 'u' models are "
f"trained with https://github.com/ultralytics/ultralytics and feature improved performance vs "
f"standard YOLOv5 models trained with https://github.com/ultralytics/yolov5.\n")
return file
def check_file(file, suffix=''):
# Search/download file (if necessary) and return path
check_suffix(file, suffix) # optional
file = str(file) # convert to str()
file = str(file) # convert to string
file = check_yolov5u_filename(file) # yolov5n -> yolov5nu
if not file or ('://' not in file and Path(file).is_file()): # exists ('://' check required in Windows Python<3.10)
return file
elif file.lower().startswith(('https://', 'http://', 'rtsp://', 'rtmp://')): # download

View File

@ -1,7 +1,6 @@
# Ultralytics YOLO 🚀, GPL-3.0 license
import contextlib
import re
import subprocess
from itertools import repeat
from multiprocessing.pool import ThreadPool
@ -111,6 +110,7 @@ def safe_download(url,
def attempt_download_asset(file, repo='ultralytics/assets', release='v0.0.0'):
# Attempt file download from GitHub release assets if not found locally. release = 'latest', 'v6.2', etc.
from ultralytics.yolo.utils import SETTINGS
from ultralytics.yolo.utils.checks import check_yolov5u_filename
def github_assets(repository, version='latest'):
# Return GitHub repo tag and assets (i.e. ['yolov8n.pt', 'yolov8s.pt', ...])
@ -121,15 +121,7 @@ def attempt_download_asset(file, repo='ultralytics/assets', release='v0.0.0'):
# YOLOv3/5u updates
file = str(file)
if 'yolov3' in file or 'yolov5' in file and 'u' not in file:
original_file = file
file = re.sub(r"(.*yolov5([nsmlx]))\.pt", "\\1u.pt", file) # i.e. yolov5n.pt -> yolov5nu.pt
file = re.sub(r"(.*yolov3(|-tiny|-spp))\.pt", "\\1u.pt", file) # i.e. yolov3-spp.pt -> yolov3-sppu.pt
if file != original_file:
LOGGER.info(f"PRO TIP 💡 Replace 'model={original_file}' with new 'model={file}'.\nYOLOv5 'u' models are "
f"trained with https://github.com/ultralytics/ultralytics and feature improved performance vs "
f"standard YOLOv5 models trained with https://github.com/ultralytics/yolov5.\n")
file = check_yolov5u_filename(file)
file = Path(file.strip().replace("'", ''))
if file.exists():
return str(file)

View File

@ -23,7 +23,7 @@ class ClassificationPredictor(BasePredictor):
results = []
for i, pred in enumerate(preds):
shape = orig_img[i].shape if isinstance(orig_img, list) else orig_img.shape
results.append(Results(probs=pred.softmax(0), orig_shape=shape[:2]))
results.append(Results(probs=pred, orig_shape=shape[:2]))
return results
def write_results(self, idx, results, batch):