From 9d27e7ada4fd09bcb8e1fee6bf7143e6b6e1833e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 17 Aug 2023 18:19:05 +0200 Subject: [PATCH] New `ASSETS` and trackers GMC cleanup (#4425) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/reference/data/converter.md | 4 + docs/reference/utils/files.md | 4 - docs/reference/utils/ops.md | 4 - examples/YOLOv8-ONNXRuntime/main.py | 9 +- examples/YOLOv8-OpenCV-ONNX-Python/main.py | 4 +- tests/test_cli.py | 12 +-- tests/test_engine.py | 11 ++- tests/test_python.py | 97 +++++++++++++++++++-- ultralytics/cfg/__init__.py | 5 +- ultralytics/cfg/trackers/botsort.yaml | 2 +- ultralytics/data/converter.py | 32 ++++++- ultralytics/data/loaders.py | 4 +- ultralytics/engine/model.py | 7 +- ultralytics/engine/predictor.py | 2 +- ultralytics/models/rtdetr/train.py | 2 +- ultralytics/models/yolo/classify/predict.py | 5 +- ultralytics/models/yolo/detect/predict.py | 5 +- ultralytics/models/yolo/detect/train.py | 2 +- ultralytics/models/yolo/detect/val.py | 6 +- ultralytics/models/yolo/pose/predict.py | 5 +- ultralytics/models/yolo/segment/predict.py | 5 +- ultralytics/models/yolo/segment/train.py | 2 +- ultralytics/models/yolo/segment/val.py | 2 +- ultralytics/nn/tasks.py | 14 --- ultralytics/trackers/bot_sort.py | 3 +- ultralytics/trackers/utils/gmc.py | 50 ++--------- ultralytics/utils/__init__.py | 15 ++-- ultralytics/utils/benchmarks.py | 29 +++--- ultralytics/utils/checks.py | 11 +-- ultralytics/utils/files.py | 12 +-- ultralytics/utils/ops.py | 21 ----- ultralytics/utils/torch_utils.py | 37 +++++--- 32 files changed, 222 insertions(+), 201 deletions(-) diff --git a/docs/reference/data/converter.md b/docs/reference/data/converter.md index 6b4c3f7..2664419 100644 --- a/docs/reference/data/converter.md +++ b/docs/reference/data/converter.md @@ -13,6 +13,10 @@ keywords: Ultralytics, Data Converter, coco91_to_coco80_class, merge_multi_segme ## ::: ultralytics.data.converter.coco91_to_coco80_class

+--- +## ::: ultralytics.data.converter.coco80_to_coco91_class +

+ --- ## ::: ultralytics.data.converter.convert_coco

diff --git a/docs/reference/utils/files.md b/docs/reference/utils/files.md index caae190..0a65bee 100644 --- a/docs/reference/utils/files.md +++ b/docs/reference/utils/files.md @@ -36,7 +36,3 @@ keywords: Ultralytics, utility functions, file operations, working directory, fi --- ## ::: ultralytics.utils.files.get_latest_run

- ---- -## ::: ultralytics.utils.files.make_dirs -

diff --git a/docs/reference/utils/ops.md b/docs/reference/utils/ops.md index d45fd90..86bb53d 100644 --- a/docs/reference/utils/ops.md +++ b/docs/reference/utils/ops.md @@ -13,10 +13,6 @@ keywords: Ultralytics YOLO, Utility Operations, segment2box, make_divisible, cli ## ::: ultralytics.utils.ops.Profile

---- -## ::: ultralytics.utils.ops.coco80_to_coco91_class -

- --- ## ::: ultralytics.utils.ops.segment2box

diff --git a/examples/YOLOv8-ONNXRuntime/main.py b/examples/YOLOv8-ONNXRuntime/main.py index 3051c47..8d03182 100644 --- a/examples/YOLOv8-ONNXRuntime/main.py +++ b/examples/YOLOv8-ONNXRuntime/main.py @@ -5,7 +5,7 @@ import numpy as np import onnxruntime as ort import torch -from ultralytics.utils import ROOT, yaml_load +from ultralytics.utils import ASSETS, yaml_load from ultralytics.utils.checks import check_requirements, check_yaml @@ -198,17 +198,14 @@ class Yolov8: outputs = session.run(None, {model_inputs[0].name: img_data}) # Perform post-processing on the outputs to obtain output image. - output_img = self.postprocess(self.img, outputs) - - # Return the resulting output image - return output_img + return self.postprocess(self.img, outputs) # output image if __name__ == '__main__': # Create an argument parser to handle command-line arguments parser = argparse.ArgumentParser() parser.add_argument('--model', type=str, default='yolov8n.onnx', help='Input your ONNX model.') - parser.add_argument('--img', type=str, default=str(ROOT / 'assets/bus.jpg'), help='Path to input image.') + parser.add_argument('--img', type=str, default=str(ASSETS / 'bus.jpg'), help='Path to input image.') parser.add_argument('--conf-thres', type=float, default=0.5, help='Confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.5, help='NMS IoU threshold') args = parser.parse_args() diff --git a/examples/YOLOv8-OpenCV-ONNX-Python/main.py b/examples/YOLOv8-OpenCV-ONNX-Python/main.py index 3cecaca..76802f0 100644 --- a/examples/YOLOv8-OpenCV-ONNX-Python/main.py +++ b/examples/YOLOv8-OpenCV-ONNX-Python/main.py @@ -3,7 +3,7 @@ import argparse import cv2.dnn import numpy as np -from ultralytics.utils import ROOT, yaml_load +from ultralytics.utils import ASSETS, yaml_load from ultralytics.utils.checks import check_yaml CLASSES = yaml_load(check_yaml('coco128.yaml'))['names'] @@ -75,6 +75,6 @@ def main(onnx_model, input_image): if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--model', default='yolov8n.onnx', help='Input your onnx model.') - parser.add_argument('--img', default=str(ROOT / 'assets/bus.jpg'), help='Path to input image.') + parser.add_argument('--img', default=str(ASSETS / 'bus.jpg'), help='Path to input image.') args = parser.parse_args() main(args.model, args.img) diff --git a/tests/test_cli.py b/tests/test_cli.py index 327ff83..91dc169 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -5,7 +5,7 @@ from pathlib import Path import pytest -from ultralytics.utils import ROOT, SETTINGS +from ultralytics.utils import ASSETS, SETTINGS WEIGHTS_DIR = Path(SETTINGS['weights_dir']) TASK_ARGS = [ @@ -40,12 +40,12 @@ def test_train(task, model, data): @pytest.mark.parametrize('task,model,data', TASK_ARGS) def test_val(task, model, data): - run(f'yolo val {task} model={WEIGHTS_DIR / model}.pt data={data} imgsz=32') + run(f'yolo val {task} model={WEIGHTS_DIR / model}.pt data={data} imgsz=32 save_txt save_json') @pytest.mark.parametrize('task,model,data', TASK_ARGS) def test_predict(task, model, data): - run(f"yolo predict model={WEIGHTS_DIR / model}.pt source={ROOT / 'assets'} imgsz=32 save save_crop save_txt") + run(f'yolo predict model={WEIGHTS_DIR / model}.pt source={ASSETS} imgsz=32 save save_crop save_txt') @pytest.mark.parametrize('model,format', EXPORT_ARGS) @@ -56,11 +56,11 @@ def test_export(model, format): def test_rtdetr(task='detect', model='yolov8n-rtdetr.yaml', data='coco8.yaml'): # Warning: MUST use imgsz=640 run(f'yolo train {task} model={model} data={data} imgsz=640 epochs=1 cache=disk') - run(f"yolo predict {task} model={model} source={ROOT / 'assets/bus.jpg'} imgsz=640 save save_crop save_txt") + run(f"yolo predict {task} model={model} source={ASSETS / 'bus.jpg'} imgsz=640 save save_crop save_txt") def test_fastsam(task='segment', model=WEIGHTS_DIR / 'FastSAM-s.pt', data='coco8-seg.yaml'): - source = ROOT / 'assets/bus.jpg' + source = ASSETS / 'bus.jpg' run(f'yolo segment val {task} model={model} data={data} imgsz=32') run(f'yolo segment predict model={model} source={source} imgsz=32 save save_crop save_txt') @@ -98,7 +98,7 @@ def test_mobilesam(): model = SAM(WEIGHTS_DIR / 'mobile_sam.pt') # Source - source = ROOT / 'assets/zidane.jpg' + source = ASSETS / 'zidane.jpg' # Predict a segment based on a point prompt model.predict(source, points=[900, 370], labels=[1]) diff --git a/tests/test_engine.py b/tests/test_engine.py index 7734269..8a21cc9 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -6,14 +6,13 @@ from ultralytics import YOLO from ultralytics.cfg import get_cfg from ultralytics.engine.exporter import Exporter from ultralytics.models.yolo import classify, detect, segment -from ultralytics.utils import DEFAULT_CFG, ROOT, SETTINGS +from ultralytics.utils import ASSETS, DEFAULT_CFG, SETTINGS CFG_DET = 'yolov8n.yaml' CFG_SEG = 'yolov8n-seg.yaml' CFG_CLS = 'yolov8n-cls.yaml' # or 'squeezenet1_0' CFG = get_cfg(DEFAULT_CFG) MODEL = Path(SETTINGS['weights_dir']) / 'yolov8n' -SOURCE = ROOT / 'assets' def test_func(*args): # noqa @@ -25,7 +24,7 @@ def test_export(): exporter.add_callback('on_export_start', test_func) assert test_func in exporter.callbacks['on_export_start'], 'callback test failed' f = exporter(model=YOLO(CFG_DET).model) - YOLO(f)(SOURCE) # exported model inference + YOLO(f)(ASSETS) # exported model inference def test_detect(): @@ -49,7 +48,7 @@ def test_detect(): pred = detect.DetectionPredictor(overrides={'imgsz': [64, 64]}) pred.add_callback('on_predict_start', test_func) assert test_func in pred.callbacks['on_predict_start'], 'callback test failed' - result = pred(source=SOURCE, model=f'{MODEL}.pt') + result = pred(source=ASSETS, model=f'{MODEL}.pt') assert len(result), 'predictor test failed' overrides['resume'] = trainer.last @@ -85,7 +84,7 @@ def test_segment(): pred = segment.SegmentationPredictor(overrides={'imgsz': [64, 64]}) pred.add_callback('on_predict_start', test_func) assert test_func in pred.callbacks['on_predict_start'], 'callback test failed' - result = pred(source=SOURCE, model=f'{MODEL}-seg.pt') + result = pred(source=ASSETS, model=f'{MODEL}-seg.pt') assert len(result), 'predictor test failed' # Test resume @@ -122,5 +121,5 @@ def test_classify(): pred = classify.ClassificationPredictor(overrides={'imgsz': [64, 64]}) pred.add_callback('on_predict_start', test_func) assert test_func in pred.callbacks['on_predict_start'], 'callback test failed' - result = pred(source=SOURCE, model=trainer.best) + result = pred(source=ASSETS, model=trainer.best) assert len(result), 'predictor test failed' diff --git a/tests/test_python.py b/tests/test_python.py index 3b92041..766c9df 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -13,14 +13,14 @@ from torchvision.transforms import ToTensor from ultralytics import RTDETR, YOLO from ultralytics.data.build import load_inference_source -from ultralytics.utils import DEFAULT_CFG, LINUX, ONLINE, ROOT, SETTINGS +from ultralytics.utils import ASSETS, DEFAULT_CFG, LINUX, ONLINE, ROOT, SETTINGS from ultralytics.utils.downloads import download from ultralytics.utils.torch_utils import TORCH_1_9 WEIGHTS_DIR = Path(SETTINGS['weights_dir']) MODEL = WEIGHTS_DIR / 'path with spaces' / 'yolov8n.pt' # test spaces in path CFG = 'yolov8n.yaml' -SOURCE = ROOT / 'assets/bus.jpg' +SOURCE = ASSETS / 'bus.jpg' TMP = (ROOT / '../tests/tmp').resolve() # temp directory for test files @@ -29,9 +29,14 @@ def test_model_forward(): model(SOURCE, imgsz=32, augment=True) -def test_model_info(): +def test_model_methods(): model = YOLO(MODEL) - model.info(verbose=True) + model.info(verbose=True, detailed=True) + model = model.reset_weights() + model = model.load(MODEL) + model.to('cpu') + _ = model.names + _ = model.device def test_model_fuse(): @@ -41,7 +46,7 @@ def test_model_fuse(): def test_predict_dir(): model = YOLO(MODEL) - model(source=ROOT / 'assets', imgsz=32) + model(source=ASSETS, imgsz=32) def test_predict_img(): @@ -102,11 +107,23 @@ def test_predict_grey_and_4ch(): def test_track_stream(): # Test YouTube streaming inference (short 10 frame video) with non-default ByteTrack tracker # imgsz=160 required for tracking for higher confidence and better matches + import yaml + model = YOLO(MODEL) model.predict('https://youtu.be/G17sBkb38XQ', imgsz=96) model.track('https://ultralytics.com/assets/decelera_portrait_min.mov', imgsz=160, tracker='bytetrack.yaml') model.track('https://ultralytics.com/assets/decelera_portrait_min.mov', imgsz=160, tracker='botsort.yaml') + # Test Global Motion Compensation (GMC) methods + for gmc in 'orb', 'sift', 'ecc': + with open(ROOT / 'cfg/trackers/botsort.yaml') as f: + data = yaml.safe_load(f) + tracker = TMP / f'botsort-{gmc}.yaml' + data['gmc_method'] = gmc + with open(tracker, 'w') as f: + yaml.safe_dump(data, f) + model.track('https://ultralytics.com/assets/decelera_portrait_min.mov', imgsz=160, tracker=tracker) + def test_val(): model = YOLO(MODEL) @@ -133,7 +150,7 @@ def test_export_torchscript(): def test_export_onnx(): model = YOLO(MODEL) - f = model.export(format='onnx') + f = model.export(format='onnx', dynamic=True) YOLO(f)(SOURCE) # exported model inference @@ -173,6 +190,12 @@ def test_export_paddle(enabled=False): model.export(format='paddle') +def test_export_ncnn(enabled=False): + model = YOLO(MODEL) + f = model.export(format='ncnn') + YOLO(f)(SOURCE) # exported model inference + + def test_all_model_yamls(): for m in (ROOT / 'cfg' / 'models').rglob('*.yaml'): if 'rtdetr' in m.name: @@ -251,12 +274,13 @@ def test_data_utils(): @pytest.mark.skipif(not ONLINE, reason='environment is offline') def test_data_converter(): # Test dataset converters - from ultralytics.data.converter import convert_coco + from ultralytics.data.converter import coco80_to_coco91_class, convert_coco file = 'instances_val2017.json' download(f'https://github.com/ultralytics/yolov5/releases/download/v1.0/{file}') shutil.move(file, TMP) convert_coco(labels_dir=TMP, use_segments=True, use_keypoints=False, cls91to80=True) + coco80_to_coco91_class() def test_events(): @@ -270,9 +294,64 @@ def test_events(): events(cfg) +def test_utils_init(): + from ultralytics.utils import (get_git_branch, get_git_origin_url, get_ubuntu_version, is_github_actions_ci, + is_ubuntu) + + is_ubuntu() + get_ubuntu_version() + is_github_actions_ci() + get_git_origin_url() + get_git_branch() + + def test_utils_checks(): - from ultralytics.utils.checks import check_yolov5u_filename, git_describe + from ultralytics.utils.checks import check_requirements, check_yolov5u_filename, git_describe - check_yolov5u_filename('yolov5.pt') + check_yolov5u_filename('yolov5n.pt') # check_imshow(warn=True) git_describe(ROOT) + check_requirements() # check requirements.txt + + +def test_utils_benchmarks(): + from ultralytics.utils.benchmarks import ProfileModels + + ProfileModels(['yolov8n.yaml'], imgsz=32, min_time=1, num_timed_runs=3, num_warmup_runs=1).profile() + + +def test_utils_torchutils(): + from ultralytics.nn.modules.conv import Conv + from ultralytics.utils.torch_utils import get_flops_with_torch_profiler, profile, time_sync + + x = torch.randn(1, 64, 20, 20) + m = Conv(64, 64, k=1, s=2) + + profile(x, [m], n=3) + get_flops_with_torch_profiler(m) + time_sync() + + +def test_utils_downloads(): + from ultralytics.utils.downloads import get_google_drive_file_info + + get_google_drive_file_info('https://drive.google.com/file/d/1cqT-cJgANNrhIHCrEufUYhQ4RqiWG_lJ/view?usp=drive_link') + + +def test_utils_ops(): + from ultralytics.utils.ops import make_divisible + + make_divisible(17, 8) + + +def test_utils_files(): + from ultralytics.utils.files import file_age, file_date, get_latest_run, spaces_in_path + + file_age(SOURCE) + file_date(SOURCE) + get_latest_run(ROOT / 'runs') + + path = TMP / 'path/with spaces' + path.mkdir(parents=True, exist_ok=True) + with spaces_in_path(path) as new_path: + print(new_path) diff --git a/ultralytics/cfg/__init__.py b/ultralytics/cfg/__init__.py index d189f56..5d191b3 100644 --- a/ultralytics/cfg/__init__.py +++ b/ultralytics/cfg/__init__.py @@ -9,7 +9,7 @@ from pathlib import Path from types import SimpleNamespace from typing import Dict, List, Union -from ultralytics.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_PATH, LOGGER, ROOT, SETTINGS, SETTINGS_YAML, +from ultralytics.utils import (ASSETS, DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_PATH, LOGGER, SETTINGS, SETTINGS_YAML, IterableSimpleNamespace, __version__, checks, colorstr, deprecation_warn, yaml_load, yaml_print) @@ -415,8 +415,7 @@ def entrypoint(debug=''): # Mode if mode in ('predict', 'track') and 'source' not in overrides: - overrides['source'] = DEFAULT_CFG.source or ROOT / 'assets' if (ROOT / 'assets').exists() \ - else 'https://ultralytics.com/images/bus.jpg' + overrides['source'] = DEFAULT_CFG.source or ASSETS LOGGER.warning(f"WARNING ⚠️ 'source' is missing. Using default 'source={overrides['source']}'.") elif mode in ('train', 'val'): if 'data' not in overrides: diff --git a/ultralytics/cfg/trackers/botsort.yaml b/ultralytics/cfg/trackers/botsort.yaml index d4947c6..cbbf348 100644 --- a/ultralytics/cfg/trackers/botsort.yaml +++ b/ultralytics/cfg/trackers/botsort.yaml @@ -11,7 +11,7 @@ match_thresh: 0.8 # threshold for matching tracks # mot20: False # for tracker evaluation(not used for now) # BoT-SORT settings -cmc_method: sparseOptFlow # method of global motion compensation +gmc_method: sparseOptFlow # method of global motion compensation # ReID model related thresh (not supported yet) proximity_thresh: 0.5 appearance_thresh: 0.25 diff --git a/ultralytics/data/converter.py b/ultralytics/data/converter.py index 218b5fa..3138102 100644 --- a/ultralytics/data/converter.py +++ b/ultralytics/data/converter.py @@ -1,6 +1,7 @@ # Ultralytics YOLO 🚀, AGPL-3.0 license import json +import shutil from collections import defaultdict from pathlib import Path @@ -9,7 +10,6 @@ import numpy as np from tqdm import tqdm from ultralytics.utils.checks import check_requirements -from ultralytics.utils.files import make_dirs def coco91_to_coco80_class(): @@ -27,6 +27,27 @@ def coco91_to_coco80_class(): None, 73, 74, 75, 76, 77, 78, 79, None] +def coco80_to_coco91_class(): # + """ + Converts 80-index (val2014) to 91-index (paper). + For details see https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/. + + Example: + ```python + import numpy as np + + a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n') + b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n') + x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco + x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet + ``` + """ + return [ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, + 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, + 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90] + + def convert_coco(labels_dir='../coco/annotations/', use_segments=False, use_keypoints=False, cls91to80=True): """Converts COCO dataset annotations to a format suitable for training YOLOv5 models. @@ -47,7 +68,14 @@ def convert_coco(labels_dir='../coco/annotations/', use_segments=False, use_keyp Generates output files in the specified output directory. """ - save_dir = make_dirs('yolo_labels') # output directory + # Create dataset directory + save_dir = Path('yolo_labels') + if save_dir.exists(): + shutil.rmtree(save_dir) # delete dir + for p in save_dir / 'labels', save_dir / 'images': + p.mkdir(parents=True, exist_ok=True) # make dir + + # Convert classes coco80 = coco91_to_coco80_class() # Import json diff --git a/ultralytics/data/loaders.py b/ultralytics/data/loaders.py index 13e8f47..1534a50 100644 --- a/ultralytics/data/loaders.py +++ b/ultralytics/data/loaders.py @@ -16,7 +16,7 @@ import torch from PIL import Image from ultralytics.data.utils import IMG_FORMATS, VID_FORMATS -from ultralytics.utils import LOGGER, ROOT, is_colab, is_kaggle, ops +from ultralytics.utils import ASSETS, LOGGER, is_colab, is_kaggle, ops from ultralytics.utils.checks import check_requirements @@ -403,7 +403,7 @@ def get_best_youtube_url(url, use_pafy=False): if __name__ == '__main__': - img = cv2.imread(str(ROOT / 'assets/bus.jpg')) + img = cv2.imread(str(ASSETS / 'bus.jpg')) dataset = LoadPilAndNumpy(im0=img) for d in dataset: print(d[0]) diff --git a/ultralytics/engine/model.py b/ultralytics/engine/model.py index 435ca85..721979b 100644 --- a/ultralytics/engine/model.py +++ b/ultralytics/engine/model.py @@ -9,8 +9,8 @@ from ultralytics.cfg import get_cfg from ultralytics.engine.exporter import Exporter from ultralytics.hub.utils import HUB_WEB_ROOT from ultralytics.nn.tasks import attempt_load_one_weight, guess_model_task, nn, yaml_model_load -from ultralytics.utils import (DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, ROOT, callbacks, emojis, - is_git_dir, yaml_load) +from ultralytics.utils import (ASSETS, DEFAULT_CFG, DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, RANK, callbacks, emojis, + yaml_load) from ultralytics.utils.checks import check_file, check_imgsz, check_pip_update_available, check_yaml from ultralytics.utils.downloads import GITHUB_ASSET_STEMS from ultralytics.utils.torch_utils import smart_inference_mode @@ -218,7 +218,7 @@ class Model: (List[ultralytics.engine.results.Results]): The prediction results. """ if source is None: - source = ROOT / 'assets' if is_git_dir() else 'https://ultralytics.com/images/bus.jpg' + source = ASSETS 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')) @@ -390,6 +390,7 @@ class Model: """ self._check_is_pytorch_model() self.model.to(device) + return self def tune(self, *args, **kwargs): """ diff --git a/ultralytics/engine/predictor.py b/ultralytics/engine/predictor.py index 6406d18..ebd51bb 100644 --- a/ultralytics/engine/predictor.py +++ b/ultralytics/engine/predictor.py @@ -47,7 +47,7 @@ STREAM_WARNING = """ WARNING ⚠️ stream/video/webcam/dir predict source will accumulate results in RAM unless `stream=True` is passed, causing potential out-of-memory errors for large sources or long-running streams/videos. - Usage: + Example: results = model(source=..., stream=True) # generator of Results objects for r in results: boxes = r.boxes # Boxes object for bbox outputs diff --git a/ultralytics/models/rtdetr/train.py b/ultralytics/models/rtdetr/train.py index dc3bc7f..a900491 100644 --- a/ultralytics/models/rtdetr/train.py +++ b/ultralytics/models/rtdetr/train.py @@ -59,7 +59,7 @@ class RTDETRTrainer(DetectionTrainer): def train(cfg=DEFAULT_CFG, use_python=False): """Train and optimize RTDETR model given training data and device.""" model = 'rtdetr-l.yaml' - data = cfg.data or 'coco128.yaml' # or yolo.ClassificationDataset("mnist") + data = cfg.data or 'coco8.yaml' # or yolo.ClassificationDataset("mnist") device = cfg.device if cfg.device is not None else '' # NOTE: F.grid_sample which is in rt-detr does not support deterministic=True diff --git a/ultralytics/models/yolo/classify/predict.py b/ultralytics/models/yolo/classify/predict.py index 5c02458..8e2f594 100644 --- a/ultralytics/models/yolo/classify/predict.py +++ b/ultralytics/models/yolo/classify/predict.py @@ -4,7 +4,7 @@ import torch from ultralytics.engine.predictor import BasePredictor from ultralytics.engine.results import Results -from ultralytics.utils import DEFAULT_CFG, ROOT +from ultralytics.utils import ASSETS, DEFAULT_CFG class ClassificationPredictor(BasePredictor): @@ -35,8 +35,7 @@ class ClassificationPredictor(BasePredictor): def predict(cfg=DEFAULT_CFG, use_python=False): """Run YOLO model predictions on input images/videos.""" model = cfg.model or 'yolov8n-cls.pt' # or "resnet18" - source = cfg.source if cfg.source is not None else ROOT / 'assets' if (ROOT / 'assets').exists() \ - else 'https://ultralytics.com/images/bus.jpg' + source = cfg.source or ASSETS args = dict(model=model, source=source) if use_python: diff --git a/ultralytics/models/yolo/detect/predict.py b/ultralytics/models/yolo/detect/predict.py index f2358c1..88b134b 100644 --- a/ultralytics/models/yolo/detect/predict.py +++ b/ultralytics/models/yolo/detect/predict.py @@ -4,7 +4,7 @@ import torch from ultralytics.engine.predictor import BasePredictor from ultralytics.engine.results import Results -from ultralytics.utils import DEFAULT_CFG, ROOT, ops +from ultralytics.utils import ASSETS, DEFAULT_CFG, ops class DetectionPredictor(BasePredictor): @@ -32,8 +32,7 @@ class DetectionPredictor(BasePredictor): def predict(cfg=DEFAULT_CFG, use_python=False): """Runs YOLO model inference on input image(s).""" model = cfg.model or 'yolov8n.pt' - source = cfg.source if cfg.source is not None else ROOT / 'assets' if (ROOT / 'assets').exists() \ - else 'https://ultralytics.com/images/bus.jpg' + source = cfg.source or ASSETS args = dict(model=model, source=source) if use_python: diff --git a/ultralytics/models/yolo/detect/train.py b/ultralytics/models/yolo/detect/train.py index 732640f..e0eeef7 100644 --- a/ultralytics/models/yolo/detect/train.py +++ b/ultralytics/models/yolo/detect/train.py @@ -107,7 +107,7 @@ class DetectionTrainer(BaseTrainer): def train(cfg=DEFAULT_CFG, use_python=False): """Train and optimize YOLO model given training data and device.""" model = cfg.model or 'yolov8n.pt' - data = cfg.data or 'coco128.yaml' # or yolo.ClassificationDataset("mnist") + data = cfg.data or 'coco8.yaml' # or yolo.ClassificationDataset("mnist") device = cfg.device if cfg.device is not None else '' args = dict(model=model, data=data, device=device) diff --git a/ultralytics/models/yolo/detect/val.py b/ultralytics/models/yolo/detect/val.py index 9eff20d..d6fb7e1 100644 --- a/ultralytics/models/yolo/detect/val.py +++ b/ultralytics/models/yolo/detect/val.py @@ -6,7 +6,7 @@ from pathlib import Path import numpy as np import torch -from ultralytics.data import build_dataloader, build_yolo_dataset +from ultralytics.data import build_dataloader, build_yolo_dataset, converter from ultralytics.engine.validator import BaseValidator from ultralytics.utils import DEFAULT_CFG, LOGGER, ops from ultralytics.utils.checks import check_requirements @@ -50,7 +50,7 @@ class DetectionValidator(BaseValidator): """Initialize evaluation metrics for YOLO.""" val = self.data.get(self.args.split, '') # validation path self.is_coco = isinstance(val, str) and 'coco' in val and val.endswith(f'{os.sep}val2017.txt') # is COCO - self.class_map = ops.coco80_to_coco91_class() if self.is_coco else list(range(1000)) + self.class_map = converter.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.names = model.names self.nc = len(model.names) @@ -259,7 +259,7 @@ class DetectionValidator(BaseValidator): def val(cfg=DEFAULT_CFG, use_python=False): """Validate trained YOLO model on validation dataset.""" model = cfg.model or 'yolov8n.pt' - data = cfg.data or 'coco128.yaml' + data = cfg.data or 'coco8.yaml' args = dict(model=model, data=data) if use_python: diff --git a/ultralytics/models/yolo/pose/predict.py b/ultralytics/models/yolo/pose/predict.py index 709e9d9..ffafadf 100644 --- a/ultralytics/models/yolo/pose/predict.py +++ b/ultralytics/models/yolo/pose/predict.py @@ -2,7 +2,7 @@ from ultralytics.engine.results import Results from ultralytics.models.yolo.detect.predict import DetectionPredictor -from ultralytics.utils import DEFAULT_CFG, LOGGER, ROOT, ops +from ultralytics.utils import ASSETS, DEFAULT_CFG, LOGGER, ops class PosePredictor(DetectionPredictor): @@ -45,8 +45,7 @@ class PosePredictor(DetectionPredictor): def predict(cfg=DEFAULT_CFG, use_python=False): """Runs YOLO to predict objects in an image or video.""" model = cfg.model or 'yolov8n-pose.pt' - source = cfg.source if cfg.source is not None else ROOT / 'assets' if (ROOT / 'assets').exists() \ - else 'https://ultralytics.com/images/bus.jpg' + source = cfg.source or ASSETS args = dict(model=model, source=source) if use_python: diff --git a/ultralytics/models/yolo/segment/predict.py b/ultralytics/models/yolo/segment/predict.py index cef9d7a..c30efe6 100644 --- a/ultralytics/models/yolo/segment/predict.py +++ b/ultralytics/models/yolo/segment/predict.py @@ -4,7 +4,7 @@ import torch from ultralytics.engine.results import Results from ultralytics.models.yolo.detect.predict import DetectionPredictor -from ultralytics.utils import DEFAULT_CFG, ROOT, ops +from ultralytics.utils import ASSETS, DEFAULT_CFG, ops class SegmentationPredictor(DetectionPredictor): @@ -47,8 +47,7 @@ class SegmentationPredictor(DetectionPredictor): def predict(cfg=DEFAULT_CFG, use_python=False): """Runs YOLO object detection on an image or video source.""" model = cfg.model or 'yolov8n-seg.pt' - source = cfg.source if cfg.source is not None else ROOT / 'assets' if (ROOT / 'assets').exists() \ - else 'https://ultralytics.com/images/bus.jpg' + source = cfg.source or ASSETS args = dict(model=model, source=source) if use_python: diff --git a/ultralytics/models/yolo/segment/train.py b/ultralytics/models/yolo/segment/train.py index f3694b9..e61d7fd 100644 --- a/ultralytics/models/yolo/segment/train.py +++ b/ultralytics/models/yolo/segment/train.py @@ -49,7 +49,7 @@ class SegmentationTrainer(yolo.detect.DetectionTrainer): def train(cfg=DEFAULT_CFG, use_python=False): """Train a YOLO segmentation model based on passed arguments.""" model = cfg.model or 'yolov8n-seg.pt' - data = cfg.data or 'coco128-seg.yaml' # or yolo.ClassificationDataset("mnist") + data = cfg.data or 'coco8-seg.yaml' device = cfg.device if cfg.device is not None else '' args = dict(model=model, data=data, device=device) diff --git a/ultralytics/models/yolo/segment/val.py b/ultralytics/models/yolo/segment/val.py index 39db266..6cabcaf 100644 --- a/ultralytics/models/yolo/segment/val.py +++ b/ultralytics/models/yolo/segment/val.py @@ -236,7 +236,7 @@ class SegmentationValidator(DetectionValidator): def val(cfg=DEFAULT_CFG, use_python=False): """Validate trained YOLO model on validation data.""" model = cfg.model or 'yolov8n-seg.pt' - data = cfg.data or 'coco128-seg.yaml' + data = cfg.data or 'coco8-seg.yaml' args = dict(model=model, data=data) if use_python: diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index dd601af..140222d 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -303,13 +303,6 @@ class SegmentationModel(DetectionModel): def init_criterion(self): return v8SegmentationLoss(self) - def _predict_augment(self, x): - """Perform augmentations on input image x and return augmented inference.""" - LOGGER.warning( - f'WARNING ⚠️ {self.__class__.__name__} has not supported augment inference yet! Now using single-scale inference instead.' - ) - return self._predict_once(x) - class PoseModel(DetectionModel): """YOLOv8 pose model.""" @@ -326,13 +319,6 @@ class PoseModel(DetectionModel): def init_criterion(self): return v8PoseLoss(self) - def _predict_augment(self, x): - """Perform augmentations on input image x and return augmented inference.""" - LOGGER.warning( - f'WARNING ⚠️ {self.__class__.__name__} has not supported augment inference yet! Now using single-scale inference instead.' - ) - return self._predict_once(x) - class ClassificationModel(BaseModel): """YOLOv8 classification model.""" diff --git a/ultralytics/trackers/bot_sort.py b/ultralytics/trackers/bot_sort.py index 168b942..d42d46e 100644 --- a/ultralytics/trackers/bot_sort.py +++ b/ultralytics/trackers/bot_sort.py @@ -110,8 +110,7 @@ class BOTSORT(BYTETracker): if args.with_reid: # Haven't supported BoT-SORT(reid) yet self.encoder = None - # self.gmc = GMC(method=args.cmc_method, verbose=[args.name, args.ablation]) - self.gmc = GMC(method=args.cmc_method) + self.gmc = GMC(method=args.gmc_method) def get_kalmanfilter(self): """Returns an instance of KalmanFilterXYWH for object tracking.""" diff --git a/ultralytics/trackers/utils/gmc.py b/ultralytics/trackers/utils/gmc.py index 0529be4..4d91df4 100644 --- a/ultralytics/trackers/utils/gmc.py +++ b/ultralytics/trackers/utils/gmc.py @@ -10,7 +10,7 @@ from ultralytics.utils import LOGGER class GMC: - def __init__(self, method='sparseOptFlow', downscale=2, verbose=None): + def __init__(self, method='sparseOptFlow', downscale=2): """Initialize a video tracker with specified parameters.""" super().__init__() @@ -40,28 +40,11 @@ class GMC: blockSize=3, useHarrisDetector=False, k=0.04) - # self.gmc_file = open('GMC_results.txt', 'w') - - elif self.method in ['file', 'files']: - seqName = verbose[0] - ablation = verbose[1] - if ablation: - filePath = r'tracker/GMC_files/MOT17_ablation' - else: - filePath = r'tracker/GMC_files/MOTChallenge' - - if '-FRCNN' in seqName: - seqName = seqName[:-6] - elif '-DPM' in seqName or '-SDP' in seqName: - seqName = seqName[:-4] - self.gmcFile = open(f'{filePath}/GMC-{seqName}.txt') - - if self.gmcFile is None: - raise ValueError(f'Error: Unable to open GMC file in directory:{filePath}') - elif self.method in ['none', 'None']: - self.method = 'none' + + elif self.method in ['none', 'None', None]: + self.method = None else: - raise ValueError(f'Error: Unknown CMC method:{method}') + raise ValueError(f'Error: Unknown GMC method:{method}') self.prevFrame = None self.prevKeyPoints = None @@ -77,10 +60,6 @@ class GMC: return self.applyEcc(raw_frame, detections) elif self.method == 'sparseOptFlow': return self.applySparseOptFlow(raw_frame, detections) - elif self.method == 'file': - return self.applyFile(raw_frame, detections) - elif self.method == 'none': - return np.eye(2, 3) else: return np.eye(2, 3) @@ -244,7 +223,6 @@ class GMC: def applySparseOptFlow(self, raw_frame, detections=None): """Initialize.""" - # t0 = time.time() height, width, _ = raw_frame.shape frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2GRAY) H = np.eye(2, 3) @@ -298,22 +276,4 @@ class GMC: self.prevFrame = frame.copy() self.prevKeyPoints = copy.copy(keypoints) - # gmc_line = str(1000 * (time.time() - t0)) + "\t" + str(H[0, 0]) + "\t" + str(H[0, 1]) + "\t" + str( - # H[0, 2]) + "\t" + str(H[1, 0]) + "\t" + str(H[1, 1]) + "\t" + str(H[1, 2]) + "\n" - # self.gmc_file.write(gmc_line) - - return H - - def applyFile(self, raw_frame, detections=None): - """Return the homography matrix based on the GCPs in the next line of the input GMC file.""" - line = self.gmcFile.readline() - tokens = line.split('\t') - H = np.eye(2, 3, dtype=np.float_) - H[0, 0] = float(tokens[1]) - H[0, 1] = float(tokens[2]) - H[0, 2] = float(tokens[3]) - H[1, 0] = float(tokens[4]) - H[1, 1] = float(tokens[5]) - H[1, 2] = float(tokens[6]) - return H diff --git a/ultralytics/utils/__init__.py b/ultralytics/utils/__init__.py index c1eb585..e8a912e 100644 --- a/ultralytics/utils/__init__.py +++ b/ultralytics/utils/__init__.py @@ -30,6 +30,7 @@ LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable # Other Constants FILE = Path(__file__).resolve() ROOT = FILE.parents[1] # YOLO +ASSETS = ROOT / 'assets' # default images DEFAULT_CFG_PATH = ROOT / 'cfg/default.yaml' NUM_THREADS = min(8, max(1, os.cpu_count() - 1)) # number of YOLOv5 multiprocessing threads AUTOINSTALL = str(os.getenv('YOLO_AUTOINSTALL', True)).lower() == 'true' # global auto-install mode @@ -260,11 +261,15 @@ class ThreadingLocked: Attributes: lock (threading.Lock): A lock object used to manage access to the decorated function. - Usage: + Example: + ```python + from ultralytics.utils import ThreadingLocked + @ThreadingLocked() def my_function(): # Your code here pass + ``` """ def __init__(self): @@ -518,7 +523,6 @@ def get_git_dir(): for d in Path(__file__).parents: if (d / '.git').is_dir(): return d - return None # no .git dir found def get_git_origin_url(): @@ -526,13 +530,12 @@ def get_git_origin_url(): Retrieves the origin URL of a git repository. Returns: - (str | None): The origin URL of the git repository. + (str | None): The origin URL of the git repository or None if not git directory. """ if is_git_dir(): with contextlib.suppress(subprocess.CalledProcessError): origin = subprocess.check_output(['git', 'config', '--get', 'remote.origin.url']) return origin.decode().strip() - return None # if not git dir or on error def get_git_branch(): @@ -540,13 +543,12 @@ def get_git_branch(): Returns the current git branch name. If not in a git repository, returns None. Returns: - (str | None): The current git branch name. + (str | None): The current git branch name or None if not a git directory. """ if is_git_dir(): with contextlib.suppress(subprocess.CalledProcessError): origin = subprocess.check_output(['git', 'rev-parse', '--abbrev-ref', 'HEAD']) return origin.decode().strip() - return None # if not git dir or on error def get_default_args(func): @@ -572,7 +574,6 @@ def get_ubuntu_version(): with contextlib.suppress(FileNotFoundError, AttributeError): with open('/etc/os-release') as f: return re.search(r'VERSION_ID="(\d+\.\d+)"', f.read())[1] - return None def get_user_config_dir(sub_dir='Ultralytics'): diff --git a/ultralytics/utils/benchmarks.py b/ultralytics/utils/benchmarks.py index 669f596..71fa57a 100644 --- a/ultralytics/utils/benchmarks.py +++ b/ultralytics/utils/benchmarks.py @@ -37,9 +37,8 @@ from tqdm import tqdm from ultralytics import YOLO from ultralytics.cfg import TASK2DATA, TASK2METRIC from ultralytics.engine.exporter import export_formats -from ultralytics.utils import LINUX, LOGGER, MACOS, ROOT, SETTINGS +from ultralytics.utils import ASSETS, LINUX, LOGGER, MACOS, SETTINGS from ultralytics.utils.checks import check_requirements, check_yolo -from ultralytics.utils.downloads import download from ultralytics.utils.files import file_size from ultralytics.utils.torch_utils import select_device @@ -68,6 +67,13 @@ def benchmark(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', Returns: df (pandas.DataFrame): A pandas DataFrame with benchmark results for each format, including file size, metric, and inference time. + + Example: + ```python + from ultralytics.utils.benchmarks import benchmark + + benchmark(model='yolov8n.pt', imgsz=640) + ``` """ import pandas as pd @@ -106,9 +112,7 @@ def benchmark(model=Path(SETTINGS['weights_dir']) / 'yolov8n.pt', assert model.task != 'pose' or i != 7, 'GraphDef Pose inference is not supported' assert i not in (9, 10), 'inference not supported' # Edge TPU and TF.js are unsupported assert i != 5 or platform.system() == 'Darwin', 'inference only supported on macOS>=10.13' # CoreML - if not (ROOT / 'assets/bus.jpg').exists(): - download(url='https://ultralytics.com/images/bus.jpg', dir=ROOT / 'assets') - export.predict(ROOT / 'assets/bus.jpg', imgsz=imgsz, device=device, half=half) + export.predict(ASSETS / 'bus.jpg', imgsz=imgsz, device=device, half=half) # Validate data = data or TASK2DATA[model.task] # task to dataset, i.e. coco8.yaml for task=detect @@ -163,6 +167,13 @@ class ProfileModels: Methods: profile(): Profiles the models and prints the result. + + Example: + ```python + from ultralytics.utils.benchmarks import ProfileModels + + ProfileModels(['yolov8n.yaml', 'yolov8s.yaml'], imgsz=640).profile() + ``` """ def __init__(self, @@ -353,11 +364,3 @@ class ProfileModels: print(separator) for row in table_rows: print(row) - - -if __name__ == '__main__': - # Benchmark all export formats - benchmark() - - # Profiling models on ONNX and TensorRT - ProfileModels(['yolov8n.yaml', 'yolov8s.yaml']) diff --git a/ultralytics/utils/checks.py b/ultralytics/utils/checks.py index 390bbea..505abcd 100644 --- a/ultralytics/utils/checks.py +++ b/ultralytics/utils/checks.py @@ -20,7 +20,7 @@ import requests import torch from matplotlib import font_manager -from ultralytics.utils import (AUTOINSTALL, LOGGER, ONLINE, ROOT, USER_CONFIG_DIR, ThreadingLocked, TryExcept, +from ultralytics.utils import (ASSETS, AUTOINSTALL, LOGGER, ONLINE, ROOT, USER_CONFIG_DIR, ThreadingLocked, TryExcept, clean_url, colorstr, downloads, emojis, is_colab, is_docker, is_jupyter, is_kaggle, is_online, is_pip_package, url2file) @@ -460,8 +460,7 @@ def check_amp(model): del m return a.shape == b.shape and torch.allclose(a, b.float(), atol=0.5) # close to 0.5 absolute tolerance - f = ROOT / 'assets/bus.jpg' # image to check - im = f if f.exists() else 'https://ultralytics.com/images/bus.jpg' if ONLINE else np.ones((640, 640, 3)) + im = ASSETS / 'bus.jpg' # image to check prefix = colorstr('AMP: ') LOGGER.info(f'{prefix}running Automatic Mixed Precision (AMP) checks with YOLOv8n...') warning_msg = "Setting 'amp=True'. If you experience zero-mAP or NaN losses you can disable AMP with amp=False." @@ -484,11 +483,9 @@ def check_amp(model): def git_describe(path=ROOT): # path must be a directory """Return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe.""" - try: - assert (Path(path) / '.git').is_dir() + with contextlib.suppress(Exception): return subprocess.check_output(f'git -C {path} describe --tags --long --always', shell=True).decode()[:-1] - except AssertionError: - return '' + return '' def print_args(args: Optional[dict] = None, show_file=True, show_func=False): diff --git a/ultralytics/utils/files.py b/ultralytics/utils/files.py index 8771ea1..0102c4b 100644 --- a/ultralytics/utils/files.py +++ b/ultralytics/utils/files.py @@ -42,6 +42,8 @@ def spaces_in_path(path): Example: ```python + with ultralytics.utils.files import spaces_in_path + with spaces_in_path('/path/with spaces') as new_path: # your code here ``` @@ -143,13 +145,3 @@ def get_latest_run(search_dir='.'): """Return path to most recent 'last.pt' in /runs (i.e. to --resume from).""" last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True) return max(last_list, key=os.path.getctime) if last_list else '' - - -def make_dirs(dir='new_dir/'): - """Create directories.""" - dir = Path(dir) - if dir.exists(): - shutil.rmtree(dir) # delete dir - for p in dir, dir / 'labels', dir / 'images': - p.mkdir(parents=True, exist_ok=True) # make dir - return dir diff --git a/ultralytics/utils/ops.py b/ultralytics/utils/ops.py index 6f11d55..1042e3c 100644 --- a/ultralytics/utils/ops.py +++ b/ultralytics/utils/ops.py @@ -55,27 +55,6 @@ class Profile(contextlib.ContextDecorator): return time.time() -def coco80_to_coco91_class(): # - """ - Converts 80-index (val2014) to 91-index (paper). - For details see https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/. - - Example: - ```python - import numpy as np - - a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n') - b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n') - x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco - x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet - ``` - """ - return [ - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, - 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, - 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90] - - def segment2box(segment, width=640, height=640): """ Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy) diff --git a/ultralytics/utils/torch_utils.py b/ultralytics/utils/torch_utils.py index 4cebcb0..def7442 100644 --- a/ultralytics/utils/torch_utils.py +++ b/ultralytics/utils/torch_utils.py @@ -239,16 +239,18 @@ def get_flops(model, imgsz=640): def get_flops_with_torch_profiler(model, imgsz=640): """Compute model FLOPs (thop alternative).""" - model = de_parallel(model) - p = next(model.parameters()) - stride = (max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32) * 2 # max stride - im = torch.zeros((1, p.shape[1], stride, stride), device=p.device) # input image in BCHW format - with torch.profiler.profile(with_flops=True) as prof: - model(im) - flops = sum(x.flops for x in prof.key_averages()) / 1E9 - imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz] # expand if int/float - flops = flops * imgsz[0] / stride * imgsz[1] / stride # 640x640 GFLOPs - return flops + if TORCH_2_0: + model = de_parallel(model) + p = next(model.parameters()) + stride = (max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32) * 2 # max stride + im = torch.zeros((1, p.shape[1], stride, stride), device=p.device) # input image in BCHW format + with torch.profiler.profile(with_flops=True) as prof: + model(im) + flops = sum(x.flops for x in prof.key_averages()) / 1E9 + imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz] # expand if int/float + flops = flops * imgsz[0] / stride * imgsz[1] / stride # 640x640 GFLOPs + return flops + return 0 def initialize_weights(model): @@ -384,11 +386,14 @@ def strip_optimizer(f: Union[str, Path] = 'best.pt', s: str = '') -> None: Returns: None - Usage: + Example: + ```python from pathlib import Path from ultralytics.utils.torch_utils import strip_optimizer - for f in Path('/Users/glennjocher/Downloads/weights').rglob('*.pt'): + + for f in Path('path/to/weights').rglob('*.pt'): strip_optimizer(f) + ``` """ # Use dill (if exists) to serialize the lambda functions where pickle does not do this try: @@ -421,13 +426,17 @@ def strip_optimizer(f: Union[str, Path] = 'best.pt', s: str = '') -> None: def profile(input, ops, n=10, device=None): """ - YOLOv8 speed/memory/FLOPs profiler + Ultralytics speed, memory and FLOPs profiler. + + Example: + ```python + from ultralytics.utils.torch_utils import profile - Usage: input = torch.randn(16, 3, 640, 640) m1 = lambda x: x * torch.sigmoid(x) m2 = nn.SiLU() profile(input, [m1, m2], n=100) # profile over 100 iterations + ``` """ results = [] if not isinstance(device, torch.device):