xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-24 06:07:03 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Release 8.0.4 fixes (#256)

Browse Source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
Co-authored-by: TechieG <35962141+gokulnath30@users.noreply.github.com>
Co-authored-by: Parthiban Marimuthu <66585214+partheee@users.noreply.github.com>
This commit is contained in:
Ayush Chaurasia 2023-01-11 23:09:52 +05:30 committed by GitHub
parent f5dfd5be8b
commit 216cf2ddb6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
18 changed files with 96 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
README.md
View File

@ -99,8 +99,8 @@ results = model("https://ultralytics.com/images/bus.jpg") # predict on an image
success = YOLO("yolov8n.pt").export(format="onnx") # export a model to ONNX format success = YOLO("yolov8n.pt").export(format="onnx") # export a model to ONNX format
``` ```
[Models](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/yolo/v8/models) download automatically from the latest [Models](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/models) download automatically from the latest
Ultralytics [release](https://github.com/ultralytics/ultralytics/releases). Ultralytics [release](https://github.com/ultralytics/assets/releases).
### Known Issues / TODOs ### Known Issues / TODOs
@ -116,18 +116,18 @@ We are still working on several parts of YOLOv8! We aim to have these completed
All YOLOv8 pretrained models are available here. Detection and Segmentation models are pretrained on the COCO dataset, while Classification models are pretrained on the ImageNet dataset. All YOLOv8 pretrained models are available here. Detection and Segmentation models are pretrained on the COCO dataset, while Classification models are pretrained on the ImageNet dataset.
[Models](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/yolo/v8/models) download automatically from the latest [Models](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/models) download automatically from the latest
Ultralytics [release](https://github.com/ultralytics/ultralytics/releases) on first use. Ultralytics [release](https://github.com/ultralytics/ultralytics/releases) on first use.
<details open><summary>Detection</summary> <details open><summary>Detection</summary>
| Model | size<br><sup>(pixels) | mAP<sup>val<br>50-95 | Speed<br><sup>CPU<br>(ms) | Speed<br><sup>T4 GPU<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) | | Model | size<br><sup>(pixels) | mAP<sup>val<br>50-95 | Speed<br><sup>CPU<br>(ms) | Speed<br><sup>T4 GPU<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) |
| ----------------------------------------------------------------------------------------- | --------------------- | -------------------- | ------------------------- | ---------------------------- | ------------------ | ----------------- | | ------------------------------------------------------------------------------------ | --------------------- | -------------------- | ------------------------- | ---------------------------- | ------------------ | ----------------- |
| [YOLOv8n](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8n.pt) | 640 | 37.3 | - | - | 3.2 | 8.7 | | [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt) | 640 | 37.3 | - | - | 3.2 | 8.7 |
| [YOLOv8s](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8s.pt) | 640 | 44.9 | - | - | 11.2 | 28.6 | | [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s.pt) | 640 | 44.9 | - | - | 11.2 | 28.6 |
| [YOLOv8m](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8m.pt) | 640 | 50.2 | - | - | 25.9 | 78.9 | | [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m.pt) | 640 | 50.2 | - | - | 25.9 | 78.9 |
| [YOLOv8l](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8l.pt) | 640 | 52.9 | - | - | 43.7 | 165.2 | | [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l.pt) | 640 | 52.9 | - | - | 43.7 | 165.2 |
| [YOLOv8x](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8x.pt) | 640 | 53.9 | - | - | 68.2 | 257.8 | | [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x.pt) | 640 | 53.9 | - | - | 68.2 | 257.8 |
- **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset. - **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset.
<br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0` <br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0`
@ -138,13 +138,13 @@ Ultralytics [release](https://github.com/ultralytics/ultralytics/releases) on fi
<details><summary>Segmentation</summary> <details><summary>Segmentation</summary>
| Model | size<br><sup>(pixels) | mAP<sup>box<br>50-95 | mAP<sup>mask<br>50-95 | Speed<br><sup>CPU<br>(ms) | Speed<br><sup>T4 GPU<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) | | Model | size<br><sup>(pixels) | mAP<sup>box<br>50-95 | mAP<sup>mask<br>50-95 | Speed<br><sup>CPU<br>(ms) | Speed<br><sup>T4 GPU<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) |
| --------------------------------------------------------------------------------------------- | --------------------- | -------------------- | --------------------- | ------------------------- | ---------------------------- | ------------------ | ----------------- | | ---------------------------------------------------------------------------------------- | --------------------- | -------------------- | --------------------- | ------------------------- | ---------------------------- | ------------------ | ----------------- |
| [YOLOv8n](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8n-seg.pt) | 640 | 36.7 | 30.5 | - | - | 3.4 | 12.6 | | [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-seg.pt) | 640 | 36.7 | 30.5 | - | - | 3.4 | 12.6 |
| [YOLOv8s](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8s-seg.pt) | 640 | 44.6 | 36.8 | - | - | 11.8 | 42.6 | | [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-seg.pt) | 640 | 44.6 | 36.8 | - | - | 11.8 | 42.6 |
| [YOLOv8m](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8m-seg.pt) | 640 | 49.9 | 40.8 | - | - | 27.3 | 110.2 | | [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m-seg.pt) | 640 | 49.9 | 40.8 | - | - | 27.3 | 110.2 |
| [YOLOv8l](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8l-seg.pt) | 640 | 52.3 | 42.6 | - | - | 46.0 | 220.5 | | [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l-seg.pt) | 640 | 52.3 | 42.6 | - | - | 46.0 | 220.5 |
| [YOLOv8x](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8x-seg.pt) | 640 | 53.4 | 43.4 | - | - | 71.8 | 344.1 | | [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-seg.pt) | 640 | 53.4 | 43.4 | - | - | 71.8 | 344.1 |
- **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset. - **mAP<sup>val</sup>** values are for single-model single-scale on [COCO val2017](http://cocodataset.org) dataset.
<br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0` <br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0`
@ -155,13 +155,13 @@ Ultralytics [release](https://github.com/ultralytics/ultralytics/releases) on fi
<details><summary>Classification</summary> <details><summary>Classification</summary>
| Model | size<br><sup>(pixels) | acc<br><sup>top1 | acc<br><sup>top5 | Speed<br><sup>CPU<br>(ms) | Speed<br><sup>T4 GPU<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) at 640 | | Model | size<br><sup>(pixels) | acc<br><sup>top1 | acc<br><sup>top5 | Speed<br><sup>CPU<br>(ms) | Speed<br><sup>T4 GPU<br>(ms) | params<br><sup>(M) | FLOPs<br><sup>(B) at 640 |
| --------------------------------------------------------------------------------------------- | --------------------- | ---------------- | ---------------- | ------------------------- | ---------------------------- | ------------------ | ------------------------ | | ---------------------------------------------------------------------------------------- | --------------------- | ---------------- | ---------------- | ------------------------- | ---------------------------- | ------------------ | ------------------------ |
| [YOLOv8n](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8n-cls.pt) | 224 | 66.6 | 87.0 | - | - | 2.7 | 4.3 | | [YOLOv8n](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n-cls.pt) | 224 | 66.6 | 87.0 | - | - | 2.7 | 4.3 |
| [YOLOv8s](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8s-cls.pt) | 224 | 72.3 | 91.1 | - | - | 6.4 | 13.5 | | [YOLOv8s](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-cls.pt) | 224 | 72.3 | 91.1 | - | - | 6.4 | 13.5 |
| [YOLOv8m](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8m-cls.pt) | 224 | 76.4 | 93.2 | - | - | 17.0 | 42.7 | | [YOLOv8m](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m-cls.pt) | 224 | 76.4 | 93.2 | - | - | 17.0 | 42.7 |
| [YOLOv8l](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8l-cls.pt) | 224 | 78.0 | 94.1 | - | - | 37.5 | 99.7 | | [YOLOv8l](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8l-cls.pt) | 224 | 78.0 | 94.1 | - | - | 37.5 | 99.7 |
| [YOLOv8x](https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8x-cls.pt) | 224 | 78.4 | 94.3 | - | - | 57.4 | 154.8 | | [YOLOv8x](https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8x-cls.pt) | 224 | 78.4 | 94.3 | - | - | 57.4 | 154.8 |
- **mAP<sup>val</sup>** values are for single-model single-scale on [ImageNet](https://www.image-net.org/) dataset. - **mAP<sup>val</sup>** values are for single-model single-scale on [ImageNet](https://www.image-net.org/) dataset.
<br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0` <br>Reproduce by `yolo mode=val task=detect data=coco.yaml device=0`

4
README.zh-CN.md
View File

@ -95,7 +95,7 @@ results = model("https://ultralytics.com/images/bus.jpg") # 预测图像
success = YOLO("yolov8n.pt").export(format="onnx") # 将模型导出为 ONNX 格式 success = YOLO("yolov8n.pt").export(format="onnx") # 将模型导出为 ONNX 格式
``` ```
[模型](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/yolo/v8/models) 会从 Ultralytics [发布页](https://github.com/ultralytics/ultralytics/releases) 自动下载。 [模型](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/models) 会从 Ultralytics [发布页](https://github.com/ultralytics/ultralytics/releases) 自动下载。
### 已知问题 / 待办事项 ### 已知问题 / 待办事项
@ -111,7 +111,7 @@ success = YOLO("yolov8n.pt").export(format="onnx") # 将模型导出为 ONNX
所有 YOLOv8 的预训练模型都可以在这里找到。目标检测和分割模型是在 COCO 数据集上预训练的，而分类模型是在 ImageNet 数据集上预训练的。 所有 YOLOv8 的预训练模型都可以在这里找到。目标检测和分割模型是在 COCO 数据集上预训练的，而分类模型是在 ImageNet 数据集上预训练的。
第一次使用时，[模型](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/yolo/v8/models) 会从 Ultralytics [发布页](https://github.com/ultralytics/ultralytics/releases) 自动下载。 第一次使用时，[模型](https://github.com/ultralytics/ultralytics/tree/main/ultralytics/models) 会从 Ultralytics [发布页](https://github.com/ultralytics/ultralytics/releases) 自动下载。
<details open><summary>目标检测</summary> <details open><summary>目标检测</summary>

6
docs/cli.md
View File

@ -55,16 +55,16 @@ You can override config file entirely by passing a new file. You can create a co
```bash ```bash
yolo task=init yolo task=init
``` ```
You can then use special `--cfg name.yaml` command to pass the new config file You can then use `cfg=name.yaml` command to pass the new config file
```bash ```bash
yolo task=detect mode=train {++ --cfg default.yaml ++} yolo cfg=default.yaml
``` ```
??? example ??? example
=== "Command" === "Command"
``` ```
yolo task=init yolo task=init
yolo task=detect mode=train --cfg default.yaml yolo cfg=default.yaml
``` ```
=== "Result" === "Result"
TODO: add terminal output TODO: add terminal output

2
tests/test_cli.py
View File

@ -23,7 +23,7 @@ def test_train_seg():
def test_train_cls(): def test_train_cls():
os.system(f'yolo mode=train task=classify model={CFG}-cls.yaml data=imagenette160 imgsz=32 epochs=1') os.system(f'yolo mode=train task=classify model={CFG}-cls.yaml data=mnist160 imgsz=32 epochs=1')
# Val checks ----------------------------------------------------------------------------------------------------------- # Val checks -----------------------------------------------------------------------------------------------------------

28
tests/test_engine.py
View File

@ -26,8 +26,10 @@ def test_detect():
# predictor # predictor
pred = detect.DetectionPredictor(overrides={"imgsz": [640, 640]}) pred = detect.DetectionPredictor(overrides={"imgsz": [640, 640]})
p = pred(source=SOURCE, model="yolov8n.pt") i = 0
assert len(p) == 2, "predictor test failed" for _ in pred(source=SOURCE, model="yolov8n.pt"):
i += 1
assert i == 2, "predictor test failed"
overrides["resume"] = trainer.last overrides["resume"] = trainer.last
trainer = detect.DetectionTrainer(overrides=overrides) trainer = detect.DetectionTrainer(overrides=overrides)
@ -57,8 +59,10 @@ def test_segment():
# predictor # predictor
pred = segment.SegmentationPredictor(overrides={"imgsz": [640, 640]}) pred = segment.SegmentationPredictor(overrides={"imgsz": [640, 640]})
p = pred(source=SOURCE, model="yolov8n-seg.pt") i = 0
assert len(p) == 2, "predictor test failed" for _ in pred(source=SOURCE, model="yolov8n-seg.pt"):
i += 1
assert i == 2, "predictor test failed"
# test resume # test resume
overrides["resume"] = trainer.last overrides["resume"] = trainer.last
@ -73,14 +77,8 @@ def test_segment():
def test_classify(): def test_classify():
overrides = { overrides = {"data": "mnist160", "model": "yolov8n-cls.yaml", "imgsz": 32, "epochs": 1, "batch": 64, "save": False}
"data": "imagenette160", CFG.data = "mnist160"
"model": "yolov8n-cls.yaml",
"imgsz": 32,
"epochs": 1,
"batch": 64,
"save": False}
CFG.data = "imagenette160"
CFG.imgsz = 32 CFG.imgsz = 32
CFG.batch = 64 CFG.batch = 64
# YOLO(CFG_SEG).train(**overrides) # This works # YOLO(CFG_SEG).train(**overrides) # This works
@ -95,5 +93,7 @@ def test_classify():
# predictor # predictor
pred = classify.ClassificationPredictor(overrides={"imgsz": [640, 640]}) pred = classify.ClassificationPredictor(overrides={"imgsz": [640, 640]})
p = pred(source=SOURCE, model=trained_model) i = 0
assert len(p) == 2, "Predictor test failed!" for _ in pred(source=SOURCE, model=trained_model):
i += 1
assert i == 2, "predictor test failed"

2
tests/test_python.py
View File

@ -32,7 +32,7 @@ def test_model_fuse():
def test_predict_dir(): def test_predict_dir():
model = YOLO(MODEL) model = YOLO(MODEL)
model.predict(source=ROOT / "assets") model.predict(source=ROOT / "assets", return_outputs=False)
def test_val(): def test_val():

1
ultralytics/nn/autobackend.py
View File

@ -56,6 +56,7 @@ class AutoBackend(nn.Module):
fp16 &= pt or jit or onnx or engine or nn_module # FP16 fp16 &= pt or jit or onnx or engine or nn_module # FP16
nhwc = coreml or saved_model or pb or tflite or edgetpu # BHWC formats (vs torch BCWH) nhwc = coreml or saved_model or pb or tflite or edgetpu # BHWC formats (vs torch BCWH)
stride = 32 # default stride stride = 32 # default stride
model = None # TODO: resolves ONNX inference, verify effect on other backends
cuda = torch.cuda.is_available() and device.type != 'cpu' # use CUDA cuda = torch.cuda.is_available() and device.type != 'cpu' # use CUDA
if not (pt or triton or nn_module): if not (pt or triton or nn_module):
w = attempt_download(w) # download if not local w = attempt_download(w) # download if not local

6
ultralytics/yolo/cli.py
View File

@ -6,6 +6,7 @@ from pathlib import Path
import hydra import hydra
from ultralytics import hub, yolo from ultralytics import hub, yolo
from ultralytics.yolo.configs import get_config
from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, colorstr from ultralytics.yolo.utils import DEFAULT_CONFIG, LOGGER, colorstr
DIR = Path(__file__).parent DIR = Path(__file__).parent
@ -20,6 +21,9 @@ def cli(cfg):
cfg (DictConfig): Configuration for the task and mode. cfg (DictConfig): Configuration for the task and mode.
""" """
# LOGGER.info(f"{colorstr(f'Ultralytics YOLO v{ultralytics.__version__}')}") # LOGGER.info(f"{colorstr(f'Ultralytics YOLO v{ultralytics.__version__}')}")
if cfg.cfg:
LOGGER.info(f"Overriding default config with {cfg.cfg}")
cfg = get_config(cfg.cfg)
task, mode = cfg.task.lower(), cfg.mode.lower() task, mode = cfg.task.lower(), cfg.mode.lower()
# Special case for initializing the configuration # Special case for initializing the configuration
@ -28,7 +32,7 @@ def cli(cfg):
LOGGER.info(f""" LOGGER.info(f"""
{colorstr("YOLO:")} configuration saved to {Path.cwd() / DEFAULT_CONFIG.name}. {colorstr("YOLO:")} configuration saved to {Path.cwd() / DEFAULT_CONFIG.name}.
To run experiments using custom configuration: To run experiments using custom configuration:
yolo task='task' mode='mode' --config-name config_file.yaml yolo cfg=config_file.yaml
""") """)
return return

1
ultralytics/yolo/configs/default.yaml
View File

@ -101,6 +101,7 @@ mixup: 0.0 # image mixup (probability)
copy_paste: 0.0 # segment copy-paste (probability) copy_paste: 0.0 # segment copy-paste (probability)
# Hydra configs -------------------------------------------------------------------------------------------------------- # Hydra configs --------------------------------------------------------------------------------------------------------
cfg: null # for overriding defaults.yaml
hydra: hydra:
output_subdir: null # disable hydra directory creation output_subdir: null # disable hydra directory creation
run: run:

10
ultralytics/yolo/engine/model.py
View File

@ -111,7 +111,7 @@ class YOLO:
self.model.fuse() self.model.fuse()
@smart_inference_mode() @smart_inference_mode()
def predict(self, source, **kwargs): def predict(self, source, return_outputs=True, **kwargs):
""" """
Visualize prediction. Visualize prediction.
@ -127,8 +127,8 @@ class YOLO:
predictor = self.PredictorClass(overrides=overrides) predictor = self.PredictorClass(overrides=overrides)
predictor.args.imgsz = check_imgsz(predictor.args.imgsz, min_dim=2) # check image size predictor.args.imgsz = check_imgsz(predictor.args.imgsz, min_dim=2) # check image size
predictor.setup(model=self.model, source=source) predictor.setup(model=self.model, source=source, return_outputs=return_outputs)
return predictor() return predictor() if return_outputs else predictor.predict_cli()
@smart_inference_mode() @smart_inference_mode()
def val(self, data=None, **kwargs): def val(self, data=None, **kwargs):
@ -212,10 +212,12 @@ class YOLO:
@staticmethod @staticmethod
def _reset_ckpt_args(args): def _reset_ckpt_args(args):
args.pop("device", None)
args.pop("project", None) args.pop("project", None)
args.pop("name", None) args.pop("name", None)
args.pop("batch", None) args.pop("batch", None)
args.pop("epochs", None) args.pop("epochs", None)
args.pop("cache", None) args.pop("cache", None)
args.pop("save_json", None) args.pop("save_json", None)
# set device to '' to prevent from auto DDP usage
args["device"] = ''

19
ultralytics/yolo/engine/predictor.py
View File

@ -89,6 +89,7 @@ class BasePredictor:
self.vid_path, self.vid_writer = None, None self.vid_path, self.vid_writer = None, None
self.annotator = None self.annotator = None
self.data_path = None self.data_path = None
self.output = dict()
self.callbacks = defaultdict(list, {k: [v] for k, v in callbacks.default_callbacks.items()}) # add callbacks self.callbacks = defaultdict(list, {k: [v] for k, v in callbacks.default_callbacks.items()}) # add callbacks
callbacks.add_integration_callbacks(self) callbacks.add_integration_callbacks(self)
@ -104,7 +105,7 @@ class BasePredictor:
def postprocess(self, preds, img, orig_img): def postprocess(self, preds, img, orig_img):
return preds return preds
def setup(self, source=None, model=None): def setup(self, source=None, model=None, return_outputs=True):
# source # source
source = str(source if source is not None else self.args.source) source = str(source if source is not None else self.args.source)
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
@ -155,16 +156,16 @@ class BasePredictor:
self.imgsz = imgsz self.imgsz = imgsz
self.done_setup = True self.done_setup = True
self.device = device self.device = device
self.return_outputs = return_outputs
return model return model
@smart_inference_mode() @smart_inference_mode()
def __call__(self, source=None, model=None): def __call__(self, source=None, model=None, return_outputs=True):
self.run_callbacks("on_predict_start") self.run_callbacks("on_predict_start")
model = self.model if self.done_setup else self.setup(source, model) model = self.model if self.done_setup else self.setup(source, model, return_outputs)
model.eval() model.eval()
self.seen, self.windows, self.dt = 0, [], (ops.Profile(), ops.Profile(), ops.Profile()) self.seen, self.windows, self.dt = 0, [], (ops.Profile(), ops.Profile(), ops.Profile())
self.all_outputs = []
for batch in self.dataset: for batch in self.dataset:
self.run_callbacks("on_predict_batch_start") self.run_callbacks("on_predict_batch_start")
path, im, im0s, vid_cap, s = batch path, im, im0s, vid_cap, s = batch
@ -194,6 +195,10 @@ class BasePredictor:
if self.args.save: if self.args.save:
self.save_preds(vid_cap, i, str(self.save_dir / p.name)) self.save_preds(vid_cap, i, str(self.save_dir / p.name))
if self.return_outputs:
yield self.output
self.output.clear()
# Print time (inference-only) # Print time (inference-only)
LOGGER.info(f"{s}{'' if len(preds) else '(no detections), '}{self.dt[1].dt * 1E3:.1f}ms") LOGGER.info(f"{s}{'' if len(preds) else '(no detections), '}{self.dt[1].dt * 1E3:.1f}ms")
@ -209,7 +214,11 @@ class BasePredictor:
LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}{s}") LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}{s}")
self.run_callbacks("on_predict_end") self.run_callbacks("on_predict_end")
return self.all_outputs
def predict_cli(self, source=None, model=None, return_outputs=False):
# as __call__ is a genertor now so have to treat it like a genertor
for _ in (self.__call__(source, model, return_outputs)):
pass
def show(self, p): def show(self, p):
im0 = self.annotator.result() im0 = self.annotator.result()

2
ultralytics/yolo/utils/torch_utils.py
View File

@ -70,7 +70,7 @@ def select_device(device='', batch_size=0, newline=False):
elif device: # non-cpu device requested elif device: # non-cpu device requested
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable - must be before assert is_available() os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable - must be before assert is_available()
assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \ assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(',', '')), \
f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)" f"Invalid CUDA 'device={device}' requested, use 'device=cpu' or pass valid CUDA device(s)"
if not cpu and not mps and torch.cuda.is_available(): # prefer GPU if available if not cpu and not mps and torch.cuda.is_available(): # prefer GPU if available
devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7 devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7

5
ultralytics/yolo/v8/classify/predict.py
View File

@ -39,7 +39,8 @@ class ClassificationPredictor(BasePredictor):
self.annotator = self.get_annotator(im0) self.annotator = self.get_annotator(im0)
prob = preds[idx].softmax(0) prob = preds[idx].softmax(0)
self.all_outputs.append(prob) if self.return_outputs:
self.output["prob"] = prob.cpu().numpy()
# Print results # Print results
top5i = prob.argsort(0, descending=True)[:5].tolist() # top 5 indices top5i = prob.argsort(0, descending=True)[:5].tolist() # top 5 indices
log_string += f"{', '.join(f'{self.model.names[j]} {prob[j]:.2f}' for j in top5i)}, " log_string += f"{', '.join(f'{self.model.names[j]} {prob[j]:.2f}' for j in top5i)}, "
@ -62,7 +63,7 @@ def predict(cfg):
cfg.source = cfg.source if cfg.source is not None else ROOT / "assets" cfg.source = cfg.source if cfg.source is not None else ROOT / "assets"
predictor = ClassificationPredictor(cfg) predictor = ClassificationPredictor(cfg)
predictor() predictor.predict_cli()
if __name__ == "__main__": if __name__ == "__main__":

1
ultralytics/yolo/v8/classify/train.py
View File

@ -143,6 +143,7 @@ def train(cfg):
cfg.weight_decay = 5e-5 cfg.weight_decay = 5e-5
cfg.label_smoothing = 0.1 cfg.label_smoothing = 0.1
cfg.warmup_epochs = 0.0 cfg.warmup_epochs = 0.0
cfg.device = cfg.device if cfg.device is not None else ''
# trainer = ClassificationTrainer(cfg) # trainer = ClassificationTrainer(cfg)
# trainer.train() # trainer.train()
from ultralytics import YOLO from ultralytics import YOLO

7
ultralytics/yolo/v8/detect/predict.py
View File

@ -53,12 +53,15 @@ class DetectionPredictor(BasePredictor):
self.annotator = self.get_annotator(im0) self.annotator = self.get_annotator(im0)
det = preds[idx] det = preds[idx]
self.all_outputs.append(det)
if len(det) == 0: if len(det) == 0:
return log_string return log_string
for c in det[:, 5].unique(): for c in det[:, 5].unique():
n = (det[:, 5] == c).sum() # detections per class n = (det[:, 5] == c).sum() # detections per class
log_string += f"{n} {self.model.names[int(c)]}{'s' * (n > 1)}, " log_string += f"{n} {self.model.names[int(c)]}{'s' * (n > 1)}, "
if self.return_outputs:
self.output["det"] = det.cpu().numpy()
# write # write
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
for *xyxy, conf, cls in reversed(det): for *xyxy, conf, cls in reversed(det):
@ -89,7 +92,7 @@ def predict(cfg):
cfg.imgsz = check_imgsz(cfg.imgsz, min_dim=2) # check image size cfg.imgsz = check_imgsz(cfg.imgsz, min_dim=2) # check image size
cfg.source = cfg.source if cfg.source is not None else ROOT / "assets" cfg.source = cfg.source if cfg.source is not None else ROOT / "assets"
predictor = DetectionPredictor(cfg) predictor = DetectionPredictor(cfg)
predictor() predictor.predict_cli()
if __name__ == "__main__": if __name__ == "__main__":

1
ultralytics/yolo/v8/detect/train.py
View File

@ -199,6 +199,7 @@ class Loss:
def train(cfg): def train(cfg):
cfg.model = cfg.model or "yolov8n.yaml" cfg.model = cfg.model or "yolov8n.yaml"
cfg.data = cfg.data or "coco128.yaml" # or yolo.ClassificationDataset("mnist") cfg.data = cfg.data or "coco128.yaml" # or yolo.ClassificationDataset("mnist")
cfg.device = cfg.device if cfg.device is not None else ''
# trainer = DetectionTrainer(cfg) # trainer = DetectionTrainer(cfg)
# trainer.train() # trainer.train()
from ultralytics import YOLO from ultralytics import YOLO

19
ultralytics/yolo/v8/segment/predict.py
View File

@ -58,10 +58,10 @@ class SegmentationPredictor(DetectionPredictor):
return log_string return log_string
# Segments # Segments
mask = masks[idx] mask = masks[idx]
if self.args.save_txt: if self.args.save_txt or self.return_outputs:
shape = im0.shape if self.args.retina_masks else im.shape[2:]
segments = [ segments = [
ops.scale_segments(im0.shape if self.args.retina_masks else im.shape[2:], x, im0.shape, normalize=True) ops.scale_segments(shape, x, im0.shape, normalize=False) for x in reversed(ops.masks2segments(mask))]
for x in reversed(ops.masks2segments(mask))]
# Print results # Print results
for c in det[:, 5].unique(): for c in det[:, 5].unique():
@ -76,12 +76,17 @@ class SegmentationPredictor(DetectionPredictor):
255 if self.args.retina_masks else im[idx]) 255 if self.args.retina_masks else im[idx])
det = reversed(det[:, :6]) det = reversed(det[:, :6])
self.all_outputs.append([det, mask]) if self.return_outputs:
self.output["det"] = det.cpu().numpy()
self.output["segment"] = segments
# Write results # Write results
for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])): for j, (*xyxy, conf, cls) in enumerate(det):
if self.args.save_txt: # Write to file if self.args.save_txt: # Write to file
seg = segments[j].reshape(-1) # (n,2) to (n*2) seg = segments[j].copy()
seg[:, 0] /= shape[1] # width
seg[:, 1] /= shape[0] # height
seg = seg.reshape(-1) # (n,2) to (n*2)
line = (cls, *seg, conf) if self.args.save_conf else (cls, *seg) # label format line = (cls, *seg, conf) if self.args.save_conf else (cls, *seg) # label format
with open(f'{self.txt_path}.txt', 'a') as f: with open(f'{self.txt_path}.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n') f.write(('%g ' * len(line)).rstrip() % line + '\n')
@ -106,7 +111,7 @@ def predict(cfg):
cfg.source = cfg.source if cfg.source is not None else ROOT / "assets" cfg.source = cfg.source if cfg.source is not None else ROOT / "assets"
predictor = SegmentationPredictor(cfg) predictor = SegmentationPredictor(cfg)
predictor() predictor.predict_cli()
if __name__ == "__main__": if __name__ == "__main__":

1
ultralytics/yolo/v8/segment/train.py
View File

@ -144,6 +144,7 @@ class SegLoss(Loss):
def train(cfg): def train(cfg):
cfg.model = cfg.model or "yolov8n-seg.yaml" cfg.model = cfg.model or "yolov8n-seg.yaml"
cfg.data = cfg.data or "coco128-seg.yaml" # or yolo.ClassificationDataset("mnist") cfg.data = cfg.data or "coco128-seg.yaml" # or yolo.ClassificationDataset("mnist")
cfg.device = cfg.device if cfg.device is not None else ''
# trainer = SegmentationTrainer(cfg) # trainer = SegmentationTrainer(cfg)
# trainer.train() # trainer.train()
from ultralytics import YOLO from ultralytics import YOLO