diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index e3b2bea..89b996d 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -17,7 +17,7 @@ jobs: strategy: fail-fast: false matrix: - os: [ubuntu-latest, windows-latest, macos-latest] + os: [ubuntu-latest] python-version: ['3.10'] # requires python<=3.9 model: [yolov8n] steps: diff --git a/docs/predict.md b/docs/predict.md index be3c3a7..966af3f 100644 --- a/docs/predict.md +++ b/docs/predict.md @@ -32,10 +32,11 @@ predictor's call method. Results object consists of these component objects: -- `Results.boxes` : `Boxes` object with properties and methods for manipulating bboxes -- `Results.masks` : `Masks` object used to index masks or to get segment coordinates. -- `Results.probs` : `torch.Tensor` containing the class probabilities/logits. -- `Results.orig_shape` : `tuple` containing the original image size as (height, width). +- `Results.boxes`: `Boxes` object with properties and methods for manipulating bboxes +- `Results.masks`: `Masks` object used to index masks or to get segment coordinates. +- `Results.probs`: `torch.Tensor` containing the class probabilities/logits. +- `Results.orig_img`: Original image loaded in memory. +- `Results.path`: `Path` containing the path to input image Each result is composed of torch.Tensor by default, in which you can easily use following functionality: @@ -94,18 +95,18 @@ results[0].probs # cls prob, (num_class, ) Class reference documentation for `Results` module and its components can be found [here](reference/results.md) -## Visualizing results +## Plotting results -You can use `visualize()` function of `Result` object to get a visualization. It plots all components(boxes, masks, +You can use `plot()` function of `Result` object to plot results on in image object. It plots all components(boxes, masks, classification logits, etc) found in the results object ```python - res = model(img) - res_plotted = res[0].visualize() - cv2.imshow("result", res_plotted) +res = model(img) +res_plotted = res[0].plot() +cv2.imshow("result", res_plotted) ``` -!!! example "`visualize()` arguments" +!!! example "`plot()` arguments" `show_conf (bool)`: Show confidence diff --git a/docs/tasks/tracking.md b/docs/tasks/tracking.md new file mode 100644 index 0000000..259487b --- /dev/null +++ b/docs/tasks/tracking.md @@ -0,0 +1,86 @@ +Object tracking is a task that involves identifying the location and class of objects, then assigning a unique ID to +that detection in video streams. + +The output of tracker is the same as detection with an added object ID. + +## Available Trackers + +The following tracking algorithms have been implemented and can be enabled by passing `tracker=tracker_type.yaml` + +* [BoT-SORT](https://github.com/NirAharon/BoT-SORT) - `botsort.yaml` +* [ByteTrack](https://github.com/ifzhang/ByteTrack) - `bytetrack.yaml` + +The default tracker is BoT-SORT. + +## Tracking + +Use a trained YOLOv8n/YOLOv8n-seg model to run tracker on video streams. + +!!! example "" + + === "Python" + + ```python + from ultralytics import YOLO + + # Load a model + model = YOLO("yolov8n.pt") # load an official detection model + model = YOLO("yolov8n-seg.pt") # load an official segmentation model + model = YOLO("path/to/best.pt") # load a custom model + + # Track with the model + results = model.track(source="https://youtu.be/Zgi9g1ksQHc", show=True) + results = model.track(source="https://youtu.be/Zgi9g1ksQHc", show=True, tracker="bytetrack.yaml") + ``` + === "CLI" + + ```bash + yolo track model=yolov8n.pt source="https://youtu.be/Zgi9g1ksQHc" # official detection model + yolo track model=yolov8n-seg.pt source=... # official segmentation model + yolo track model=path/to/best.pt source=... # custom model + yolo track model=path/to/best.pt tracker="bytetrack.yaml" # bytetrack tracker + + ``` + +As in the above usage, we support both the detection and segmentation models for tracking and the only thing you need to do is loading the corresponding(detection or segmentation) model. + +## Configuration +### Tracking +Tracking shares the configuration with predict, i.e `conf`, `iou`, `show`. More configurations please refer to [predict page](https://docs.ultralytics.com/cfg/#prediction). +!!! example "" + + === "Python" + + ```python + from ultralytics import YOLO + + model = YOLO("yolov8n.pt") + results = model.track(source="https://youtu.be/Zgi9g1ksQHc", conf=0.3, iou=0.5, show=True) + ``` + === "CLI" + + ```bash + yolo track model=yolov8n.pt source="https://youtu.be/Zgi9g1ksQHc" conf=0.3, iou=0.5 show + + ``` + +### Tracker +We also support using a modified tracker config file, just copy a config file i.e `custom_tracker.yaml` from [ultralytics/tracker/cfg](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/tracker/cfg) and modify any configurations(expect the `tracker_type`) you need to. +!!! example "" + + === "Python" + + ```python + from ultralytics import YOLO + + model = YOLO("yolov8n.pt") + results = model.track(source="https://youtu.be/Zgi9g1ksQHc", tracker='custom_tracker.yaml') + ``` + === "CLI" + + ```bash + yolo track model=yolov8n.pt source="https://youtu.be/Zgi9g1ksQHc" tracker='custom_tracker.yaml' + + ``` +Please refer to [ultralytics/tracker/cfg](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/tracker/cfg) page. + diff --git a/examples/README.md b/examples/README.md index c7ba2da..36b4a51 100644 --- a/examples/README.md +++ b/examples/README.md @@ -2,10 +2,11 @@ This is a list of real-world applications and walkthroughs. These can be folders ## Ultralytics YOLO example applications -| Title | Format | Contributor | -| --------------------------------------------------------------- | ------------------ | --------------------------------------------------- | -| [Yolov8/yolov5 ONNX Inference with C++](./YOLOv8-CPP-Inference) | C++/ONNX | [Justas Bartnykas](https://github.com/JustasBart) | -| [YOLOv8-OpenCV-ONNX-Python](./YOLOv8-OpenCV-ONNX-Python) | OpenCV/Python/ONNX | [Farid Inawan](https://github.com/frdteknikelektro) | +| Title | Format | Contributor | +| ------------------------------------------------------------------------ | ------------------ | --------------------------------------------------- | +| [YOLO ONNX detection Inference with C++](./YOLOv8_CPP_Inference) | C++/ONNX | [Justas Bartnykas](https://github.com/JustasBart) | +| [YOLO OpenCV ONNX detection Python](./YOLOv8-OpenCV-ONNX-Python) | OpenCV/Python/ONNX | [Farid Inawan](https://github.com/frdteknikelektro) | +| [YOLO .Net ONNX detection C#](https://www.nuget.org/packages/Yolov8.Net) | C# .Net | [Samuel Stainback](https://github.com/sstainba) | ## How can you contribute ? diff --git a/examples/tutorial.ipynb b/examples/tutorial.ipynb index b2685e4..9bf0355 100644 --- a/examples/tutorial.ipynb +++ b/examples/tutorial.ipynb @@ -123,7 +123,7 @@ "Downloading https://ultralytics.com/images/zidane.jpg to zidane.jpg...\n", "100% 165k/165k [00:00<00:00, 87.4MB/s]\n", "image 1/1 /content/zidane.jpg: 384x640 2 persons, 1 tie, 13.3ms\n", - "Speed: 0.5ms pre-process, 13.3ms inference, 43.5ms postprocess per image at shape (1, 3, 640, 640)\n", + "Speed: 0.5ms preprocess, 13.3ms inference, 43.5ms postprocess per image at shape (1, 3, 640, 640)\n", "Results saved to \u001b[1mruns/detect/predict\u001b[0m\n" ] } @@ -268,7 +268,7 @@ " scissors 128 1 1 0 0.249 0.0746\n", " teddy bear 128 21 0.877 0.333 0.591 0.394\n", " toothbrush 128 5 0.743 0.6 0.638 0.374\n", - "Speed: 2.4ms pre-process, 7.8ms inference, 0.0ms loss, 3.3ms post-process per image\n" + "Speed: 2.4ms preprocess, 7.8ms inference, 0.0ms loss, 3.3ms postprocess per image\n" ] } ] @@ -439,7 +439,7 @@ " scissors 128 1 1 0 0.142 0.0426\n", " teddy bear 128 21 0.587 0.476 0.63 0.458\n", " toothbrush 128 5 0.784 0.736 0.898 0.544\n", - "Speed: 2.0ms pre-process, 4.0ms inference, 0.0ms loss, 2.5ms post-process per image\n", + "Speed: 2.0ms preprocess, 4.0ms inference, 0.0ms loss, 2.5ms postprocess per image\n", "Results saved to \u001b[1mruns/detect/train\u001b[0m\n" ] } diff --git a/mkdocs.yml b/mkdocs.yml index eef6271..29f92ca 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -105,6 +105,7 @@ nav: - Tasks: - Detection: tasks/detection.md - Segmentation: tasks/segmentation.md + - Multi-Object Tracking: tasks/tracking.md - Classification: tasks/classification.md - Usage: - CLI: cli.md diff --git a/tests/test_python.py b/tests/test_python.py index 1ac6dc4..86425ae 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -170,15 +170,16 @@ def test_predict_callback_and_setup(): def test_result(): model = YOLO('yolov8n-seg.pt') res = model([SOURCE, SOURCE]) - res[0].numpy() res[0].cpu().numpy() - resimg = res[0].visualize(show_conf=False) - print(resimg) + res[0].plot(show_conf=False) + print(res[0].path) model = YOLO('yolov8n.pt') res = model(SOURCE) - res[0].visualize() + res[0].plot() + print(res[0].path) model = YOLO('yolov8n-cls.pt') res = model(SOURCE) - res[0].visualize() + res[0].plot() + print(res[0].path) diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index 89e4905..6102be4 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -1,6 +1,6 @@ # Ultralytics YOLO 🚀, GPL-3.0 license -__version__ = '8.0.41' +__version__ = '8.0.42' from ultralytics.yolo.engine.model import YOLO from ultralytics.yolo.utils.checks import check_yolo as checks diff --git a/ultralytics/nn/autoshape.py b/ultralytics/nn/autoshape.py index 654a688..8b1b920 100644 --- a/ultralytics/nn/autoshape.py +++ b/ultralytics/nn/autoshape.py @@ -232,6 +232,3 @@ class Detections: def __repr__(self): return f'YOLOv8 {self.__class__} instance\n' + self.__str__() - - -print('works') diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index 6a663f0..29ad140 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -381,7 +381,7 @@ def attempt_load_weights(weights, device=None, inplace=True, fuse=False): return ensemble[-1] # Return ensemble - print(f'Ensemble created with {weights}\n') + LOGGER.info(f'Ensemble created with {weights}\n') for k in 'names', 'nc', 'yaml': setattr(ensemble, k, getattr(ensemble[0], k)) ensemble.stride = ensemble[torch.argmax(torch.tensor([m.stride.max() for m in ensemble])).int()].stride diff --git a/ultralytics/tracker/README.md b/ultralytics/tracker/README.md index fcf5c64..387ca98 100644 --- a/ultralytics/tracker/README.md +++ b/ultralytics/tracker/README.md @@ -16,7 +16,7 @@ model = YOLO("yolov8n.pt") # or a segmentation model .i.e yolov8n-seg.pt model.track( source="video/streams", stream=True, - tracker="botsort.yaml/bytetrack.yaml", + tracker="botsort.yaml", # or 'bytetrack.yaml' ..., ) ``` diff --git a/ultralytics/tracker/__init__.py b/ultralytics/tracker/__init__.py index 2eb9f41..da0f3ab 100644 --- a/ultralytics/tracker/__init__.py +++ b/ultralytics/tracker/__init__.py @@ -1 +1,3 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license + from .trackers import BOTSORT, BYTETracker diff --git a/ultralytics/tracker/cfg/botsort.yaml b/ultralytics/tracker/cfg/botsort.yaml index d2fd3ca..445d1a1 100644 --- a/ultralytics/tracker/cfg/botsort.yaml +++ b/ultralytics/tracker/cfg/botsort.yaml @@ -1,3 +1,6 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license +# Default YOLO tracker settings for BoT-SORT tracker https://github.com/NirAharon/BoT-SORT + tracker_type: botsort # tracker type, ['botsort', 'bytetrack'] track_high_thresh: 0.5 # threshold for the first association track_low_thresh: 0.1 # threshold for the second association @@ -7,7 +10,7 @@ match_thresh: 0.8 # threshold for matching tracks # min_box_area: 10 # threshold for min box areas(for tracker evaluation, not used for now) # mot20: False # for tracker evaluation(not used for now) -# Botsort settings +# BoT-SORT settings cmc_method: sparseOptFlow # method of global motion compensation # ReID model related thresh (not supported yet) proximity_thresh: 0.5 diff --git a/ultralytics/tracker/cfg/bytetrack.yaml b/ultralytics/tracker/cfg/bytetrack.yaml index 9d1860f..fe9378c 100644 --- a/ultralytics/tracker/cfg/bytetrack.yaml +++ b/ultralytics/tracker/cfg/bytetrack.yaml @@ -1,3 +1,6 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license +# Default YOLO tracker settings for ByteTrack tracker https://github.com/ifzhang/ByteTrack + tracker_type: bytetrack # tracker type, ['botsort', 'bytetrack'] track_high_thresh: 0.5 # threshold for the first association track_low_thresh: 0.1 # threshold for the second association diff --git a/ultralytics/tracker/track.py b/ultralytics/tracker/track.py index be16e6b..9fd8a75 100644 --- a/ultralytics/tracker/track.py +++ b/ultralytics/tracker/track.py @@ -1,3 +1,5 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license + import torch from ultralytics.tracker import BOTSORT, BYTETracker diff --git a/ultralytics/tracker/trackers/__init__.py b/ultralytics/tracker/trackers/__init__.py index 225217c..ae99ab6 100644 --- a/ultralytics/tracker/trackers/__init__.py +++ b/ultralytics/tracker/trackers/__init__.py @@ -1,2 +1,4 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license + from .bot_sort import BOTSORT from .byte_tracker import BYTETracker diff --git a/ultralytics/tracker/trackers/basetrack.py b/ultralytics/tracker/trackers/basetrack.py index c19464a..a173f03 100644 --- a/ultralytics/tracker/trackers/basetrack.py +++ b/ultralytics/tracker/trackers/basetrack.py @@ -1,3 +1,5 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license + from collections import OrderedDict import numpy as np diff --git a/ultralytics/tracker/trackers/bot_sort.py b/ultralytics/tracker/trackers/bot_sort.py index fab20a6..d520e8e 100644 --- a/ultralytics/tracker/trackers/bot_sort.py +++ b/ultralytics/tracker/trackers/bot_sort.py @@ -1,3 +1,5 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license + from collections import deque import numpy as np @@ -97,7 +99,7 @@ class BOTSORT(BYTETracker): self.appearance_thresh = args.appearance_thresh if args.with_reid: - # haven't supported bot-sort(reid) yet + # 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) diff --git a/ultralytics/tracker/trackers/byte_tracker.py b/ultralytics/tracker/trackers/byte_tracker.py index eedc208..2d5dfd1 100644 --- a/ultralytics/tracker/trackers/byte_tracker.py +++ b/ultralytics/tracker/trackers/byte_tracker.py @@ -1,3 +1,5 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license + import numpy as np from ..utils import matching diff --git a/ultralytics/tracker/utils/gmc.py b/ultralytics/tracker/utils/gmc.py index 705d6aa..653b5e8 100644 --- a/ultralytics/tracker/utils/gmc.py +++ b/ultralytics/tracker/utils/gmc.py @@ -1,9 +1,13 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license + import copy import cv2 import matplotlib.pyplot as plt import numpy as np +from ultralytics.yolo.utils import LOGGER + class GMC: @@ -108,7 +112,7 @@ class GMC: try: (cc, H) = cv2.findTransformECC(self.prevFrame, frame, H, self.warp_mode, self.criteria, None, 1) except Exception as e: - print(f'Warning: find transform failed. Set warp as identity {e}') + LOGGER.warning(f'WARNING: find transform failed. Set warp as identity {e}') return H @@ -229,7 +233,7 @@ class GMC: H[0, 2] *= self.downscale H[1, 2] *= self.downscale else: - print('Warning: not enough matching points') + LOGGER.warning('WARNING: not enough matching points') # Store to next iteration self.prevFrame = frame.copy() @@ -288,7 +292,7 @@ class GMC: H[0, 2] *= self.downscale H[1, 2] *= self.downscale else: - print('Warning: not enough matching points') + LOGGER.warning('WARNING: not enough matching points') # Store to next iteration self.prevFrame = frame.copy() diff --git a/ultralytics/tracker/utils/kalman_filter.py b/ultralytics/tracker/utils/kalman_filter.py index d6fb42f..1d12805 100644 --- a/ultralytics/tracker/utils/kalman_filter.py +++ b/ultralytics/tracker/utils/kalman_filter.py @@ -1,3 +1,5 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license + import numpy as np import scipy.linalg @@ -234,7 +236,7 @@ class KalmanFilterXYAH: class KalmanFilterXYWH: """ - For bot-sort + For BoT-SORT A simple Kalman filter for tracking bounding boxes in image space. The 8-dimensional state space diff --git a/ultralytics/tracker/utils/matching.py b/ultralytics/tracker/utils/matching.py index e19fc3f..38840d7 100644 --- a/ultralytics/tracker/utils/matching.py +++ b/ultralytics/tracker/utils/matching.py @@ -1,3 +1,5 @@ +# Ultralytics YOLO 🚀, GPL-3.0 license + import lap import numpy as np import scipy diff --git a/ultralytics/yolo/data/dataloaders/v5loader.py b/ultralytics/yolo/data/dataloaders/v5loader.py index 613c681..b7eb7f4 100644 --- a/ultralytics/yolo/data/dataloaders/v5loader.py +++ b/ultralytics/yolo/data/dataloaders/v5loader.py @@ -1136,11 +1136,11 @@ class HUBDatasetStats(): # Save, print and return if save: stats_path = self.hub_dir / 'stats.json' - print(f'Saving {stats_path.resolve()}...') + LOGGER.info(f'Saving {stats_path.resolve()}...') with open(stats_path, 'w') as f: json.dump(self.stats, f) # save stats.json if verbose: - print(json.dumps(self.stats, indent=2, sort_keys=False)) + LOGGER.info(json.dumps(self.stats, indent=2, sort_keys=False)) return self.stats def process_images(self): @@ -1154,7 +1154,7 @@ class HUBDatasetStats(): with ThreadPool(NUM_THREADS) as pool: for _ in tqdm(pool.imap(self._hub_ops, dataset.im_files), total=total, desc=desc): pass - print(f'Done. All images saved to {self.im_dir}') + LOGGER.info(f'Done. All images saved to {self.im_dir}') return self.im_dir diff --git a/ultralytics/yolo/engine/exporter.py b/ultralytics/yolo/engine/exporter.py index cbe0046..f77ab01 100644 --- a/ultralytics/yolo/engine/exporter.py +++ b/ultralytics/yolo/engine/exporter.py @@ -75,7 +75,6 @@ from ultralytics.yolo.utils.files import file_size from ultralytics.yolo.utils.ops import Profile from ultralytics.yolo.utils.torch_utils import get_latest_opset, select_device, smart_inference_mode -CUDA = torch.cuda.is_available() ARM64 = platform.machine() in ('arm64', 'aarch64') @@ -324,7 +323,7 @@ class Exporter: # Simplify if self.args.simplify: try: - check_requirements(('onnxsim', 'onnxruntime-gpu' if CUDA else 'onnxruntime')) + check_requirements(('onnxsim', 'onnxruntime-gpu' if torch.cuda.is_available() else 'onnxruntime')) import onnxsim LOGGER.info(f'{prefix} simplifying with onnxsim {onnxsim.__version__}...') @@ -506,10 +505,12 @@ class Exporter: try: import tensorflow as tf # noqa except ImportError: - check_requirements(f"tensorflow{'-macos' if MACOS else '-aarch64' if ARM64 else '' if CUDA else '-cpu'}") + check_requirements( + f"tensorflow{'-macos' if MACOS else '-aarch64' if ARM64 else '' if torch.cuda.is_available() else '-cpu'}" + ) import tensorflow as tf # noqa check_requirements(('onnx', 'onnx2tf', 'sng4onnx', 'onnxsim', 'onnx_graphsurgeon', 'tflite_support', - 'onnxruntime-gpu' if CUDA else 'onnxruntime'), + 'onnxruntime-gpu' if torch.cuda.is_available() else 'onnxruntime'), cmds='--extra-index-url https://pypi.ngc.nvidia.com') LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index cb92c91..e468180 100644 --- a/ultralytics/yolo/engine/model.py +++ b/ultralytics/yolo/engine/model.py @@ -32,7 +32,7 @@ class YOLO: YOLO (You Only Look Once) object detection model. Args: - model (str or Path): Path to the model file to load or create. + model (str, Path): Path to the model file to load or create. type (str): Type/version of models to use. Defaults to "v8". Attributes: @@ -62,7 +62,7 @@ class YOLO: predict(source=None, stream=False, **kwargs): Perform prediction using the YOLO model. Returns: - List[ultralytics.yolo.engine.results.Results]: The prediction results. + list(ultralytics.yolo.engine.results.Results): The prediction results. """ def __init__(self, model='yolov8n.pt', type='v8') -> None: @@ -114,6 +114,7 @@ class YOLO: self.task = guess_model_task(cfg_dict) self.ModelClass, self.TrainerClass, self.ValidatorClass, self.PredictorClass = self._assign_ops_from_task() self.model = self.ModelClass(cfg_dict, verbose=verbose and RANK == -1) # initialize + self.overrides['model'] = self.cfg def _load(self, weights: str): """ @@ -204,7 +205,7 @@ class YOLO: def track(self, source=None, stream=False, **kwargs): from ultralytics.tracker.track import register_tracker register_tracker(self) - # bytetrack-based method needs low confidence predictions as input + # ByteTrack-based method needs low confidence predictions as input conf = kwargs.get('conf') or 0.1 kwargs['conf'] = conf kwargs['mode'] = 'track' diff --git a/ultralytics/yolo/engine/predictor.py b/ultralytics/yolo/engine/predictor.py index 9f4857d..132a504 100644 --- a/ultralytics/yolo/engine/predictor.py +++ b/ultralytics/yolo/engine/predictor.py @@ -92,6 +92,7 @@ class BasePredictor: self.annotator = None self.data_path = None self.source_type = None + self.batch = None self.callbacks = defaultdict(list, callbacks.default_callbacks) # add callbacks callbacks.add_integration_callbacks(self) diff --git a/ultralytics/yolo/engine/results.py b/ultralytics/yolo/engine/results.py index 702220d..9e7a3fd 100644 --- a/ultralytics/yolo/engine/results.py +++ b/ultralytics/yolo/engine/results.py @@ -28,13 +28,14 @@ class Results: """ - def __init__(self, boxes=None, masks=None, probs=None, orig_img=None, names=None) -> None: + def __init__(self, orig_img, path, names, boxes=None, masks=None, probs=None) -> None: self.orig_img = orig_img self.orig_shape = orig_img.shape[:2] self.boxes = Boxes(boxes, self.orig_shape) if boxes is not None else None # native size boxes self.masks = Masks(masks, self.orig_shape) if masks is not None else None # native size or imgsz masks self.probs = probs if probs is not None else None self.names = names + self.path = path self.comp = ['boxes', 'masks', 'probs'] def pandas(self): @@ -42,7 +43,7 @@ class Results: # TODO masks.pandas + boxes.pandas + cls.pandas def __getitem__(self, idx): - r = Results(orig_img=self.orig_img) + r = Results(orig_img=self.orig_img, path=self.path, names=self.names) for item in self.comp: if getattr(self, item) is None: continue @@ -58,7 +59,7 @@ class Results: self.probs = probs def cpu(self): - r = Results(orig_img=self.orig_img) + r = Results(orig_img=self.orig_img, path=self.path, names=self.names) for item in self.comp: if getattr(self, item) is None: continue @@ -66,7 +67,7 @@ class Results: return r def numpy(self): - r = Results(orig_img=self.orig_img) + r = Results(orig_img=self.orig_img, path=self.path, names=self.names) for item in self.comp: if getattr(self, item) is None: continue @@ -74,7 +75,7 @@ class Results: return r def cuda(self): - r = Results(orig_img=self.orig_img) + r = Results(orig_img=self.orig_img, path=self.path, names=self.names) for item in self.comp: if getattr(self, item) is None: continue @@ -82,7 +83,7 @@ class Results: return r def to(self, *args, **kwargs): - r = Results(orig_img=self.orig_img) + r = Results(orig_img=self.orig_img, path=self.path, names=self.names) for item in self.comp: if getattr(self, item) is None: continue @@ -123,7 +124,7 @@ class Results: orig_shape (tuple, optional): Original image size. """) - def visualize(self, show_conf=True, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'): + def plot(self, show_conf=True, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'): """ Plots the given result on an input RGB image. Accepts cv2(numpy) or PIL Image @@ -146,9 +147,9 @@ class Results: annotator.box_label(d.xyxy.squeeze(), label, color=colors(c, True)) if masks is not None: - im_gpu = torch.as_tensor(img, dtype=torch.float16).permute(2, 0, 1).flip(0).contiguous() - im_gpu = F.resize(im_gpu, masks.data.shape[1:]) / 255 - annotator.masks(masks.data, colors=[colors(x, True) for x in boxes.cls], im_gpu=im_gpu) + im = torch.as_tensor(img, dtype=torch.float16, device=masks.data.device).permute(2, 0, 1).flip(0) + im = F.resize(im.contiguous(), masks.data.shape[1:]) / 255 + annotator.masks(masks.data, colors=[colors(x, True) for x in boxes.cls], im_gpu=im) if logits is not None: top5i = logits.argsort(0, descending=True)[:5].tolist() # top 5 indices @@ -371,24 +372,3 @@ class Masks: Properties: segments (list): A list of segments which includes x,y,w,h,label,confidence, and mask of each detection masks. """) - - -if __name__ == '__main__': - # test examples - results = Results(boxes=torch.randn((2, 6)), masks=torch.randn((2, 160, 160)), orig_shape=[640, 640]) - results = results.cuda() - print('--cuda--pass--') - results = results.cpu() - print('--cpu--pass--') - results = results.to('cuda:0') - print('--to-cuda--pass--') - results = results.to('cpu') - print('--to-cpu--pass--') - results = results.numpy() - print('--numpy--pass--') - # box = Boxes(boxes=torch.randn((2, 6)), orig_shape=[5, 5]) - # box = box.cuda() - # box = box.cpu() - # box = box.numpy() - # for b in box: - # print(b) diff --git a/ultralytics/yolo/utils/metrics.py b/ultralytics/yolo/utils/metrics.py index de0deb0..61f1490 100644 --- a/ultralytics/yolo/utils/metrics.py +++ b/ultralytics/yolo/utils/metrics.py @@ -11,7 +11,7 @@ import numpy as np import torch import torch.nn as nn -from ultralytics.yolo.utils import TryExcept +from ultralytics.yolo.utils import LOGGER, TryExcept # boxes @@ -260,7 +260,7 @@ class ConfusionMatrix: def print(self): for i in range(self.nc + 1): - print(' '.join(map(str, self.matrix[i]))) + LOGGER.info(' '.join(map(str, self.matrix[i]))) def smooth(y, f=0.05): diff --git a/ultralytics/yolo/utils/plotting.py b/ultralytics/yolo/utils/plotting.py index ddbc830..5d5da9f 100644 --- a/ultralytics/yolo/utils/plotting.py +++ b/ultralytics/yolo/utils/plotting.py @@ -12,7 +12,7 @@ import torch from PIL import Image, ImageDraw, ImageFont from PIL import __version__ as pil_version -from ultralytics.yolo.utils import threaded +from ultralytics.yolo.utils import LOGGER, threaded from .checks import check_font, check_version, is_ascii from .files import increment_path @@ -300,7 +300,7 @@ def plot_results(file='path/to/results.csv', dir='', segment=False): # if j in [8, 9, 10]: # share train and val loss y axes # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5]) except Exception as e: - print(f'Warning: Plotting error for {f}: {e}') + LOGGER.warning(f'WARNING: Plotting error for {f}: {e}') ax[1].legend() fig.savefig(save_dir / 'results.png', dpi=200) plt.close() diff --git a/ultralytics/yolo/utils/torch_utils.py b/ultralytics/yolo/utils/torch_utils.py index 670d208..34ddf0d 100644 --- a/ultralytics/yolo/utils/torch_utils.py +++ b/ultralytics/yolo/utils/torch_utils.py @@ -167,11 +167,12 @@ def model_info(model, verbose=False, imgsz=640): n_p = get_num_params(model) n_g = get_num_gradients(model) # number gradients if verbose: - print(f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}") + LOGGER.info( + f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}") for i, (name, p) in enumerate(model.named_parameters()): name = name.replace('module_list.', '') - print('%5g %40s %9s %12g %20s %10.3g %10.3g' % - (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) + LOGGER.info('%5g %40s %9s %12g %20s %10.3g %10.3g' % + (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) flops = get_flops(model, imgsz) fused = ' (fused)' if model.is_fused() else '' @@ -362,8 +363,8 @@ def profile(input, ops, n=10, device=None): results = [] if not isinstance(device, torch.device): device = select_device(device) - print(f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}" - f"{'input':>24s}{'output':>24s}") + LOGGER.info(f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}" + f"{'input':>24s}{'output':>24s}") for x in input if isinstance(input, list) else [input]: x = x.to(device) @@ -393,10 +394,10 @@ def profile(input, ops, n=10, device=None): mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0 # (GB) s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else 'list' for x in (x, y)) # shapes p = sum(x.numel() for x in m.parameters()) if isinstance(m, nn.Module) else 0 # parameters - print(f'{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}') + LOGGER.info(f'{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}') results.append([p, flops, mem, tf, tb, s_in, s_out]) except Exception as e: - print(e) + LOGGER.info(e) results.append(None) torch.cuda.empty_cache() return results diff --git a/ultralytics/yolo/v8/classify/predict.py b/ultralytics/yolo/v8/classify/predict.py index 596190d..18c868e 100644 --- a/ultralytics/yolo/v8/classify/predict.py +++ b/ultralytics/yolo/v8/classify/predict.py @@ -22,7 +22,9 @@ class ClassificationPredictor(BasePredictor): results = [] for i, pred in enumerate(preds): orig_img = orig_img[i] if isinstance(orig_img, list) else orig_img - results.append(Results(probs=pred, orig_img=orig_img, names=self.model.names)) + path, _, _, _, _ = self.batch + img_path = path[i] if isinstance(path, list) else path + results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, probs=pred)) return results diff --git a/ultralytics/yolo/v8/detect/predict.py b/ultralytics/yolo/v8/detect/predict.py index 1f47c1b..f26a20a 100644 --- a/ultralytics/yolo/v8/detect/predict.py +++ b/ultralytics/yolo/v8/detect/predict.py @@ -32,7 +32,9 @@ class DetectionPredictor(BasePredictor): orig_img = orig_img[i] if isinstance(orig_img, list) else orig_img shape = orig_img.shape pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round() - results.append(Results(boxes=pred, orig_img=orig_img, names=self.model.names)) + path, _, _, _, _ = self.batch + img_path = path[i] if isinstance(path, list) else path + results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred)) return results def write_results(self, idx, results, batch): diff --git a/ultralytics/yolo/v8/segment/predict.py b/ultralytics/yolo/v8/segment/predict.py index 2335087..24bc9d6 100644 --- a/ultralytics/yolo/v8/segment/predict.py +++ b/ultralytics/yolo/v8/segment/predict.py @@ -24,9 +24,10 @@ class SegmentationPredictor(DetectionPredictor): for i, pred in enumerate(p): orig_img = orig_img[i] if isinstance(orig_img, list) else orig_img shape = orig_img.shape - if not len(pred): - results.append(Results(boxes=pred[:, :6], orig_img=orig_img, - names=self.model.names)) # save empty boxes + path, _, _, _, _ = self.batch + img_path = path[i] if isinstance(path, list) else path + if not len(pred): # save empty boxes + results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6])) continue if self.args.retina_masks: pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round() @@ -34,7 +35,8 @@ class SegmentationPredictor(DetectionPredictor): else: masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True) # HWC pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round() - results.append(Results(boxes=pred[:, :6], masks=masks, orig_img=orig_img, names=self.model.names)) + results.append( + Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks)) return results def write_results(self, idx, results, batch):