ultralytics 8.1.29__py3-none-any.whl → 8.3.63__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- tests/__init__.py +22 -0
- tests/conftest.py +83 -0
- tests/test_cli.py +122 -0
- tests/test_cuda.py +155 -0
- tests/test_engine.py +131 -0
- tests/test_exports.py +216 -0
- tests/test_integrations.py +150 -0
- tests/test_python.py +615 -0
- tests/test_solutions.py +94 -0
- ultralytics/__init__.py +11 -8
- ultralytics/cfg/__init__.py +569 -131
- ultralytics/cfg/datasets/Argoverse.yaml +2 -1
- ultralytics/cfg/datasets/DOTAv1.5.yaml +3 -2
- ultralytics/cfg/datasets/DOTAv1.yaml +3 -2
- ultralytics/cfg/datasets/GlobalWheat2020.yaml +3 -2
- ultralytics/cfg/datasets/ImageNet.yaml +2 -1
- ultralytics/cfg/datasets/Objects365.yaml +5 -4
- ultralytics/cfg/datasets/SKU-110K.yaml +2 -1
- ultralytics/cfg/datasets/VOC.yaml +3 -2
- ultralytics/cfg/datasets/VisDrone.yaml +6 -5
- ultralytics/cfg/datasets/african-wildlife.yaml +25 -0
- ultralytics/cfg/datasets/brain-tumor.yaml +23 -0
- ultralytics/cfg/datasets/carparts-seg.yaml +3 -2
- ultralytics/cfg/datasets/coco-pose.yaml +7 -6
- ultralytics/cfg/datasets/coco.yaml +3 -2
- ultralytics/cfg/datasets/coco128-seg.yaml +4 -3
- ultralytics/cfg/datasets/coco128.yaml +4 -3
- ultralytics/cfg/datasets/coco8-pose.yaml +3 -2
- ultralytics/cfg/datasets/coco8-seg.yaml +3 -2
- ultralytics/cfg/datasets/coco8.yaml +3 -2
- ultralytics/cfg/datasets/crack-seg.yaml +3 -2
- ultralytics/cfg/datasets/dog-pose.yaml +24 -0
- ultralytics/cfg/datasets/dota8.yaml +3 -2
- ultralytics/cfg/datasets/hand-keypoints.yaml +26 -0
- ultralytics/cfg/datasets/lvis.yaml +1236 -0
- ultralytics/cfg/datasets/medical-pills.yaml +22 -0
- ultralytics/cfg/datasets/open-images-v7.yaml +2 -1
- ultralytics/cfg/datasets/package-seg.yaml +5 -4
- ultralytics/cfg/datasets/signature.yaml +21 -0
- ultralytics/cfg/datasets/tiger-pose.yaml +3 -2
- ultralytics/cfg/datasets/xView.yaml +2 -1
- ultralytics/cfg/default.yaml +14 -11
- ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml +24 -0
- ultralytics/cfg/models/11/yolo11-cls.yaml +33 -0
- ultralytics/cfg/models/11/yolo11-obb.yaml +50 -0
- ultralytics/cfg/models/11/yolo11-pose.yaml +51 -0
- ultralytics/cfg/models/11/yolo11-seg.yaml +50 -0
- ultralytics/cfg/models/11/yolo11.yaml +50 -0
- ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +5 -2
- ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +5 -2
- ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +5 -2
- ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +5 -2
- ultralytics/cfg/models/v10/yolov10b.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10l.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10m.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10n.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10s.yaml +45 -0
- ultralytics/cfg/models/v10/yolov10x.yaml +45 -0
- ultralytics/cfg/models/v3/yolov3-spp.yaml +5 -2
- ultralytics/cfg/models/v3/yolov3-tiny.yaml +5 -2
- ultralytics/cfg/models/v3/yolov3.yaml +5 -2
- ultralytics/cfg/models/v5/yolov5-p6.yaml +5 -2
- ultralytics/cfg/models/v5/yolov5.yaml +5 -2
- ultralytics/cfg/models/v6/yolov6.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-cls.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +6 -2
- ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +6 -2
- ultralytics/cfg/models/v8/yolov8-ghost.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-obb.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-p2.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-p6.yaml +10 -7
- ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-pose.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-seg.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-world.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8-worldv2.yaml +5 -2
- ultralytics/cfg/models/v8/yolov8.yaml +5 -2
- ultralytics/cfg/models/v9/yolov9c-seg.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9c.yaml +30 -25
- ultralytics/cfg/models/v9/yolov9e-seg.yaml +64 -0
- ultralytics/cfg/models/v9/yolov9e.yaml +46 -42
- ultralytics/cfg/models/v9/yolov9m.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9s.yaml +41 -0
- ultralytics/cfg/models/v9/yolov9t.yaml +41 -0
- ultralytics/cfg/solutions/default.yaml +24 -0
- ultralytics/cfg/trackers/botsort.yaml +8 -5
- ultralytics/cfg/trackers/bytetrack.yaml +8 -5
- ultralytics/data/__init__.py +14 -3
- ultralytics/data/annotator.py +37 -15
- ultralytics/data/augment.py +1783 -289
- ultralytics/data/base.py +62 -27
- ultralytics/data/build.py +37 -8
- ultralytics/data/converter.py +196 -36
- ultralytics/data/dataset.py +233 -94
- ultralytics/data/loaders.py +199 -96
- ultralytics/data/split_dota.py +39 -29
- ultralytics/data/utils.py +111 -41
- ultralytics/engine/__init__.py +1 -1
- ultralytics/engine/exporter.py +579 -244
- ultralytics/engine/model.py +604 -252
- ultralytics/engine/predictor.py +22 -11
- ultralytics/engine/results.py +1228 -218
- ultralytics/engine/trainer.py +191 -129
- ultralytics/engine/tuner.py +18 -18
- ultralytics/engine/validator.py +18 -15
- ultralytics/hub/__init__.py +31 -13
- ultralytics/hub/auth.py +11 -7
- ultralytics/hub/google/__init__.py +159 -0
- ultralytics/hub/session.py +128 -94
- ultralytics/hub/utils.py +20 -21
- ultralytics/models/__init__.py +4 -2
- ultralytics/models/fastsam/__init__.py +2 -3
- ultralytics/models/fastsam/model.py +26 -4
- ultralytics/models/fastsam/predict.py +127 -63
- ultralytics/models/fastsam/utils.py +1 -44
- ultralytics/models/fastsam/val.py +1 -1
- ultralytics/models/nas/__init__.py +1 -1
- ultralytics/models/nas/model.py +21 -10
- ultralytics/models/nas/predict.py +3 -6
- ultralytics/models/nas/val.py +4 -4
- ultralytics/models/rtdetr/__init__.py +1 -1
- ultralytics/models/rtdetr/model.py +1 -1
- ultralytics/models/rtdetr/predict.py +6 -8
- ultralytics/models/rtdetr/train.py +6 -2
- ultralytics/models/rtdetr/val.py +3 -3
- ultralytics/models/sam/__init__.py +3 -3
- ultralytics/models/sam/amg.py +29 -23
- ultralytics/models/sam/build.py +211 -13
- ultralytics/models/sam/model.py +91 -30
- ultralytics/models/sam/modules/__init__.py +1 -1
- ultralytics/models/sam/modules/blocks.py +1129 -0
- ultralytics/models/sam/modules/decoders.py +381 -53
- ultralytics/models/sam/modules/encoders.py +515 -324
- ultralytics/models/sam/modules/memory_attention.py +237 -0
- ultralytics/models/sam/modules/sam.py +969 -21
- ultralytics/models/sam/modules/tiny_encoder.py +425 -154
- ultralytics/models/sam/modules/transformer.py +159 -60
- ultralytics/models/sam/modules/utils.py +293 -0
- ultralytics/models/sam/predict.py +1263 -132
- ultralytics/models/utils/__init__.py +1 -1
- ultralytics/models/utils/loss.py +36 -24
- ultralytics/models/utils/ops.py +3 -7
- ultralytics/models/yolo/__init__.py +3 -3
- ultralytics/models/yolo/classify/__init__.py +1 -1
- ultralytics/models/yolo/classify/predict.py +7 -8
- ultralytics/models/yolo/classify/train.py +17 -22
- ultralytics/models/yolo/classify/val.py +8 -4
- ultralytics/models/yolo/detect/__init__.py +1 -1
- ultralytics/models/yolo/detect/predict.py +3 -5
- ultralytics/models/yolo/detect/train.py +11 -4
- ultralytics/models/yolo/detect/val.py +90 -52
- ultralytics/models/yolo/model.py +14 -9
- ultralytics/models/yolo/obb/__init__.py +1 -1
- ultralytics/models/yolo/obb/predict.py +2 -2
- ultralytics/models/yolo/obb/train.py +5 -3
- ultralytics/models/yolo/obb/val.py +41 -23
- ultralytics/models/yolo/pose/__init__.py +1 -1
- ultralytics/models/yolo/pose/predict.py +3 -5
- ultralytics/models/yolo/pose/train.py +2 -2
- ultralytics/models/yolo/pose/val.py +51 -17
- ultralytics/models/yolo/segment/__init__.py +1 -1
- ultralytics/models/yolo/segment/predict.py +3 -5
- ultralytics/models/yolo/segment/train.py +2 -2
- ultralytics/models/yolo/segment/val.py +60 -19
- ultralytics/models/yolo/world/__init__.py +5 -0
- ultralytics/models/yolo/world/train.py +92 -0
- ultralytics/models/yolo/world/train_world.py +109 -0
- ultralytics/nn/__init__.py +1 -1
- ultralytics/nn/autobackend.py +228 -93
- ultralytics/nn/modules/__init__.py +39 -14
- ultralytics/nn/modules/activation.py +21 -0
- ultralytics/nn/modules/block.py +526 -66
- ultralytics/nn/modules/conv.py +24 -7
- ultralytics/nn/modules/head.py +177 -34
- ultralytics/nn/modules/transformer.py +6 -5
- ultralytics/nn/modules/utils.py +1 -2
- ultralytics/nn/tasks.py +226 -82
- ultralytics/solutions/__init__.py +30 -1
- ultralytics/solutions/ai_gym.py +96 -143
- ultralytics/solutions/analytics.py +247 -0
- ultralytics/solutions/distance_calculation.py +78 -135
- ultralytics/solutions/heatmap.py +93 -247
- ultralytics/solutions/object_counter.py +184 -259
- ultralytics/solutions/parking_management.py +246 -0
- ultralytics/solutions/queue_management.py +112 -0
- ultralytics/solutions/region_counter.py +116 -0
- ultralytics/solutions/security_alarm.py +144 -0
- ultralytics/solutions/solutions.py +178 -0
- ultralytics/solutions/speed_estimation.py +86 -174
- ultralytics/solutions/streamlit_inference.py +190 -0
- ultralytics/solutions/trackzone.py +68 -0
- ultralytics/trackers/__init__.py +1 -1
- ultralytics/trackers/basetrack.py +32 -13
- ultralytics/trackers/bot_sort.py +61 -28
- ultralytics/trackers/byte_tracker.py +83 -51
- ultralytics/trackers/track.py +21 -6
- ultralytics/trackers/utils/__init__.py +1 -1
- ultralytics/trackers/utils/gmc.py +62 -48
- ultralytics/trackers/utils/kalman_filter.py +166 -35
- ultralytics/trackers/utils/matching.py +40 -21
- ultralytics/utils/__init__.py +511 -239
- ultralytics/utils/autobatch.py +40 -22
- ultralytics/utils/benchmarks.py +266 -85
- ultralytics/utils/callbacks/__init__.py +1 -1
- ultralytics/utils/callbacks/base.py +1 -3
- ultralytics/utils/callbacks/clearml.py +7 -6
- ultralytics/utils/callbacks/comet.py +39 -17
- ultralytics/utils/callbacks/dvc.py +1 -1
- ultralytics/utils/callbacks/hub.py +16 -16
- ultralytics/utils/callbacks/mlflow.py +28 -24
- ultralytics/utils/callbacks/neptune.py +6 -2
- ultralytics/utils/callbacks/raytune.py +3 -4
- ultralytics/utils/callbacks/tensorboard.py +18 -18
- ultralytics/utils/callbacks/wb.py +27 -20
- ultralytics/utils/checks.py +172 -100
- ultralytics/utils/dist.py +2 -1
- ultralytics/utils/downloads.py +40 -34
- ultralytics/utils/errors.py +1 -1
- ultralytics/utils/files.py +72 -38
- ultralytics/utils/instance.py +41 -19
- ultralytics/utils/loss.py +83 -55
- ultralytics/utils/metrics.py +61 -56
- ultralytics/utils/ops.py +94 -89
- ultralytics/utils/patches.py +30 -14
- ultralytics/utils/plotting.py +600 -269
- ultralytics/utils/tal.py +67 -26
- ultralytics/utils/torch_utils.py +305 -112
- ultralytics/utils/triton.py +2 -1
- ultralytics/utils/tuner.py +21 -12
- ultralytics-8.3.63.dist-info/METADATA +370 -0
- ultralytics-8.3.63.dist-info/RECORD +241 -0
- {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/WHEEL +1 -1
- ultralytics/data/explorer/__init__.py +0 -5
- ultralytics/data/explorer/explorer.py +0 -472
- ultralytics/data/explorer/gui/__init__.py +0 -1
- ultralytics/data/explorer/gui/dash.py +0 -268
- ultralytics/data/explorer/utils.py +0 -166
- ultralytics/models/fastsam/prompt.py +0 -357
- ultralytics-8.1.29.dist-info/METADATA +0 -373
- ultralytics-8.1.29.dist-info/RECORD +0 -197
- {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/LICENSE +0 -0
- {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/entry_points.txt +0 -0
- {ultralytics-8.1.29.dist-info → ultralytics-8.3.63.dist-info}/top_level.txt +0 -0
ultralytics/engine/results.py
CHANGED
@@ -1,4 +1,4 @@
|
|
1
|
-
# Ultralytics
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
2
|
"""
|
3
3
|
Ultralytics Results, Boxes and Masks classes for handling inference results.
|
4
4
|
|
@@ -14,52 +14,173 @@ import torch
|
|
14
14
|
|
15
15
|
from ultralytics.data.augment import LetterBox
|
16
16
|
from ultralytics.utils import LOGGER, SimpleClass, ops
|
17
|
+
from ultralytics.utils.checks import check_requirements
|
17
18
|
from ultralytics.utils.plotting import Annotator, colors, save_one_box
|
18
19
|
from ultralytics.utils.torch_utils import smart_inference_mode
|
19
20
|
|
20
21
|
|
21
22
|
class BaseTensor(SimpleClass):
|
22
|
-
"""
|
23
|
+
"""
|
24
|
+
Base tensor class with additional methods for easy manipulation and device handling.
|
25
|
+
|
26
|
+
Attributes:
|
27
|
+
data (torch.Tensor | np.ndarray): Prediction data such as bounding boxes, masks, or keypoints.
|
28
|
+
orig_shape (Tuple[int, int]): Original shape of the image, typically in the format (height, width).
|
29
|
+
|
30
|
+
Methods:
|
31
|
+
cpu: Return a copy of the tensor stored in CPU memory.
|
32
|
+
numpy: Returns a copy of the tensor as a numpy array.
|
33
|
+
cuda: Moves the tensor to GPU memory, returning a new instance if necessary.
|
34
|
+
to: Return a copy of the tensor with the specified device and dtype.
|
35
|
+
|
36
|
+
Examples:
|
37
|
+
>>> import torch
|
38
|
+
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
39
|
+
>>> orig_shape = (720, 1280)
|
40
|
+
>>> base_tensor = BaseTensor(data, orig_shape)
|
41
|
+
>>> cpu_tensor = base_tensor.cpu()
|
42
|
+
>>> numpy_array = base_tensor.numpy()
|
43
|
+
>>> gpu_tensor = base_tensor.cuda()
|
44
|
+
"""
|
23
45
|
|
24
46
|
def __init__(self, data, orig_shape) -> None:
|
25
47
|
"""
|
26
|
-
Initialize BaseTensor with data and original shape.
|
48
|
+
Initialize BaseTensor with prediction data and the original shape of the image.
|
27
49
|
|
28
50
|
Args:
|
29
|
-
data (torch.Tensor | np.ndarray):
|
30
|
-
orig_shape (
|
51
|
+
data (torch.Tensor | np.ndarray): Prediction data such as bounding boxes, masks, or keypoints.
|
52
|
+
orig_shape (Tuple[int, int]): Original shape of the image in (height, width) format.
|
53
|
+
|
54
|
+
Examples:
|
55
|
+
>>> import torch
|
56
|
+
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
57
|
+
>>> orig_shape = (720, 1280)
|
58
|
+
>>> base_tensor = BaseTensor(data, orig_shape)
|
31
59
|
"""
|
32
|
-
assert isinstance(data, (torch.Tensor, np.ndarray))
|
60
|
+
assert isinstance(data, (torch.Tensor, np.ndarray)), "data must be torch.Tensor or np.ndarray"
|
33
61
|
self.data = data
|
34
62
|
self.orig_shape = orig_shape
|
35
63
|
|
36
64
|
@property
|
37
65
|
def shape(self):
|
38
|
-
"""
|
66
|
+
"""
|
67
|
+
Returns the shape of the underlying data tensor.
|
68
|
+
|
69
|
+
Returns:
|
70
|
+
(Tuple[int, ...]): The shape of the data tensor.
|
71
|
+
|
72
|
+
Examples:
|
73
|
+
>>> data = torch.rand(100, 4)
|
74
|
+
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
|
75
|
+
>>> print(base_tensor.shape)
|
76
|
+
(100, 4)
|
77
|
+
"""
|
39
78
|
return self.data.shape
|
40
79
|
|
41
80
|
def cpu(self):
|
42
|
-
"""
|
81
|
+
"""
|
82
|
+
Returns a copy of the tensor stored in CPU memory.
|
83
|
+
|
84
|
+
Returns:
|
85
|
+
(BaseTensor): A new BaseTensor object with the data tensor moved to CPU memory.
|
86
|
+
|
87
|
+
Examples:
|
88
|
+
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]]).cuda()
|
89
|
+
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
|
90
|
+
>>> cpu_tensor = base_tensor.cpu()
|
91
|
+
>>> isinstance(cpu_tensor, BaseTensor)
|
92
|
+
True
|
93
|
+
>>> cpu_tensor.data.device
|
94
|
+
device(type='cpu')
|
95
|
+
"""
|
43
96
|
return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.cpu(), self.orig_shape)
|
44
97
|
|
45
98
|
def numpy(self):
|
46
|
-
"""
|
99
|
+
"""
|
100
|
+
Returns a copy of the tensor as a numpy array.
|
101
|
+
|
102
|
+
Returns:
|
103
|
+
(np.ndarray): A numpy array containing the same data as the original tensor.
|
104
|
+
|
105
|
+
Examples:
|
106
|
+
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
107
|
+
>>> orig_shape = (720, 1280)
|
108
|
+
>>> base_tensor = BaseTensor(data, orig_shape)
|
109
|
+
>>> numpy_array = base_tensor.numpy()
|
110
|
+
>>> print(type(numpy_array))
|
111
|
+
<class 'numpy.ndarray'>
|
112
|
+
"""
|
47
113
|
return self if isinstance(self.data, np.ndarray) else self.__class__(self.data.numpy(), self.orig_shape)
|
48
114
|
|
49
115
|
def cuda(self):
|
50
|
-
"""
|
116
|
+
"""
|
117
|
+
Moves the tensor to GPU memory.
|
118
|
+
|
119
|
+
Returns:
|
120
|
+
(BaseTensor): A new BaseTensor instance with the data moved to GPU memory if it's not already a
|
121
|
+
numpy array, otherwise returns self.
|
122
|
+
|
123
|
+
Examples:
|
124
|
+
>>> import torch
|
125
|
+
>>> from ultralytics.engine.results import BaseTensor
|
126
|
+
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
127
|
+
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
|
128
|
+
>>> gpu_tensor = base_tensor.cuda()
|
129
|
+
>>> print(gpu_tensor.data.device)
|
130
|
+
cuda:0
|
131
|
+
"""
|
51
132
|
return self.__class__(torch.as_tensor(self.data).cuda(), self.orig_shape)
|
52
133
|
|
53
134
|
def to(self, *args, **kwargs):
|
54
|
-
"""
|
135
|
+
"""
|
136
|
+
Return a copy of the tensor with the specified device and dtype.
|
137
|
+
|
138
|
+
Args:
|
139
|
+
*args (Any): Variable length argument list to be passed to torch.Tensor.to().
|
140
|
+
**kwargs (Any): Arbitrary keyword arguments to be passed to torch.Tensor.to().
|
141
|
+
|
142
|
+
Returns:
|
143
|
+
(BaseTensor): A new BaseTensor instance with the data moved to the specified device and/or dtype.
|
144
|
+
|
145
|
+
Examples:
|
146
|
+
>>> base_tensor = BaseTensor(torch.randn(3, 4), orig_shape=(480, 640))
|
147
|
+
>>> cuda_tensor = base_tensor.to("cuda")
|
148
|
+
>>> float16_tensor = base_tensor.to(dtype=torch.float16)
|
149
|
+
"""
|
55
150
|
return self.__class__(torch.as_tensor(self.data).to(*args, **kwargs), self.orig_shape)
|
56
151
|
|
57
152
|
def __len__(self): # override len(results)
|
58
|
-
"""
|
153
|
+
"""
|
154
|
+
Returns the length of the underlying data tensor.
|
155
|
+
|
156
|
+
Returns:
|
157
|
+
(int): The number of elements in the first dimension of the data tensor.
|
158
|
+
|
159
|
+
Examples:
|
160
|
+
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
161
|
+
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
|
162
|
+
>>> len(base_tensor)
|
163
|
+
2
|
164
|
+
"""
|
59
165
|
return len(self.data)
|
60
166
|
|
61
167
|
def __getitem__(self, idx):
|
62
|
-
"""
|
168
|
+
"""
|
169
|
+
Returns a new BaseTensor instance containing the specified indexed elements of the data tensor.
|
170
|
+
|
171
|
+
Args:
|
172
|
+
idx (int | List[int] | torch.Tensor): Index or indices to select from the data tensor.
|
173
|
+
|
174
|
+
Returns:
|
175
|
+
(BaseTensor): A new BaseTensor instance containing the indexed data.
|
176
|
+
|
177
|
+
Examples:
|
178
|
+
>>> data = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
179
|
+
>>> base_tensor = BaseTensor(data, orig_shape=(720, 1280))
|
180
|
+
>>> result = base_tensor[0] # Select the first row
|
181
|
+
>>> print(result.data)
|
182
|
+
tensor([1, 2, 3])
|
183
|
+
"""
|
63
184
|
return self.__class__(self.data[idx], self.orig_shape)
|
64
185
|
|
65
186
|
|
@@ -67,46 +188,74 @@ class Results(SimpleClass):
|
|
67
188
|
"""
|
68
189
|
A class for storing and manipulating inference results.
|
69
190
|
|
191
|
+
This class encapsulates the functionality for handling detection, segmentation, pose estimation,
|
192
|
+
and classification results from YOLO models.
|
193
|
+
|
70
194
|
Attributes:
|
71
195
|
orig_img (numpy.ndarray): Original image as a numpy array.
|
72
|
-
orig_shape (
|
73
|
-
boxes (Boxes
|
74
|
-
masks (Masks
|
75
|
-
probs (Probs
|
76
|
-
keypoints (Keypoints
|
77
|
-
|
78
|
-
|
196
|
+
orig_shape (Tuple[int, int]): Original image shape in (height, width) format.
|
197
|
+
boxes (Boxes | None): Object containing detection bounding boxes.
|
198
|
+
masks (Masks | None): Object containing detection masks.
|
199
|
+
probs (Probs | None): Object containing class probabilities for classification tasks.
|
200
|
+
keypoints (Keypoints | None): Object containing detected keypoints for each object.
|
201
|
+
obb (OBB | None): Object containing oriented bounding boxes.
|
202
|
+
speed (Dict[str, float | None]): Dictionary of preprocess, inference, and postprocess speeds.
|
203
|
+
names (Dict[int, str]): Dictionary mapping class IDs to class names.
|
79
204
|
path (str): Path to the image file.
|
205
|
+
_keys (Tuple[str, ...]): Tuple of attribute names for internal use.
|
80
206
|
|
81
207
|
Methods:
|
82
|
-
update
|
83
|
-
cpu
|
84
|
-
numpy
|
85
|
-
cuda
|
86
|
-
to
|
87
|
-
new
|
88
|
-
plot
|
89
|
-
show
|
90
|
-
save
|
91
|
-
verbose
|
92
|
-
save_txt
|
93
|
-
save_crop
|
94
|
-
tojson
|
208
|
+
update: Updates object attributes with new detection results.
|
209
|
+
cpu: Returns a copy of the Results object with all tensors on CPU memory.
|
210
|
+
numpy: Returns a copy of the Results object with all tensors as numpy arrays.
|
211
|
+
cuda: Returns a copy of the Results object with all tensors on GPU memory.
|
212
|
+
to: Returns a copy of the Results object with tensors on a specified device and dtype.
|
213
|
+
new: Returns a new Results object with the same image, path, and names.
|
214
|
+
plot: Plots detection results on an input image, returning an annotated image.
|
215
|
+
show: Shows annotated results on screen.
|
216
|
+
save: Saves annotated results to file.
|
217
|
+
verbose: Returns a log string for each task, detailing detections and classifications.
|
218
|
+
save_txt: Saves detection results to a text file.
|
219
|
+
save_crop: Saves cropped detection images.
|
220
|
+
tojson: Converts detection results to JSON format.
|
221
|
+
|
222
|
+
Examples:
|
223
|
+
>>> results = model("path/to/image.jpg")
|
224
|
+
>>> for result in results:
|
225
|
+
... print(result.boxes) # Print detection boxes
|
226
|
+
... result.show() # Display the annotated image
|
227
|
+
... result.save(filename="result.jpg") # Save annotated image
|
95
228
|
"""
|
96
229
|
|
97
|
-
def __init__(
|
230
|
+
def __init__(
|
231
|
+
self, orig_img, path, names, boxes=None, masks=None, probs=None, keypoints=None, obb=None, speed=None
|
232
|
+
) -> None:
|
98
233
|
"""
|
99
|
-
Initialize the Results class.
|
234
|
+
Initialize the Results class for storing and manipulating inference results.
|
100
235
|
|
101
236
|
Args:
|
102
237
|
orig_img (numpy.ndarray): The original image as a numpy array.
|
103
238
|
path (str): The path to the image file.
|
104
|
-
names (
|
105
|
-
boxes (torch.
|
106
|
-
masks (torch.
|
107
|
-
probs (torch.
|
108
|
-
keypoints (torch.
|
109
|
-
obb (torch.
|
239
|
+
names (Dict): A dictionary of class names.
|
240
|
+
boxes (torch.Tensor | None): A 2D tensor of bounding box coordinates for each detection.
|
241
|
+
masks (torch.Tensor | None): A 3D tensor of detection masks, where each mask is a binary image.
|
242
|
+
probs (torch.Tensor | None): A 1D tensor of probabilities of each class for classification task.
|
243
|
+
keypoints (torch.Tensor | None): A 2D tensor of keypoint coordinates for each detection.
|
244
|
+
obb (torch.Tensor | None): A 2D tensor of oriented bounding box coordinates for each detection.
|
245
|
+
speed (Dict | None): A dictionary containing preprocess, inference, and postprocess speeds (ms/image).
|
246
|
+
|
247
|
+
Examples:
|
248
|
+
>>> results = model("path/to/image.jpg")
|
249
|
+
>>> result = results[0] # Get the first result
|
250
|
+
>>> boxes = result.boxes # Get the boxes for the first result
|
251
|
+
>>> masks = result.masks # Get the masks for the first result
|
252
|
+
|
253
|
+
Notes:
|
254
|
+
For the default pose model, keypoint indices for human body pose estimation are:
|
255
|
+
0: Nose, 1: Left Eye, 2: Right Eye, 3: Left Ear, 4: Right Ear
|
256
|
+
5: Left Shoulder, 6: Right Shoulder, 7: Left Elbow, 8: Right Elbow
|
257
|
+
9: Left Wrist, 10: Right Wrist, 11: Left Hip, 12: Right Hip
|
258
|
+
13: Left Knee, 14: Right Knee, 15: Left Ankle, 16: Right Ankle
|
110
259
|
"""
|
111
260
|
self.orig_img = orig_img
|
112
261
|
self.orig_shape = orig_img.shape[:2]
|
@@ -115,25 +264,66 @@ class Results(SimpleClass):
|
|
115
264
|
self.probs = Probs(probs) if probs is not None else None
|
116
265
|
self.keypoints = Keypoints(keypoints, self.orig_shape) if keypoints is not None else None
|
117
266
|
self.obb = OBB(obb, self.orig_shape) if obb is not None else None
|
118
|
-
self.speed = {"preprocess": None, "inference": None, "postprocess": None}
|
267
|
+
self.speed = speed if speed is not None else {"preprocess": None, "inference": None, "postprocess": None}
|
119
268
|
self.names = names
|
120
269
|
self.path = path
|
121
270
|
self.save_dir = None
|
122
271
|
self._keys = "boxes", "masks", "probs", "keypoints", "obb"
|
123
272
|
|
124
273
|
def __getitem__(self, idx):
|
125
|
-
"""
|
274
|
+
"""
|
275
|
+
Return a Results object for a specific index of inference results.
|
276
|
+
|
277
|
+
Args:
|
278
|
+
idx (int | slice): Index or slice to retrieve from the Results object.
|
279
|
+
|
280
|
+
Returns:
|
281
|
+
(Results): A new Results object containing the specified subset of inference results.
|
282
|
+
|
283
|
+
Examples:
|
284
|
+
>>> results = model("path/to/image.jpg") # Perform inference
|
285
|
+
>>> single_result = results[0] # Get the first result
|
286
|
+
>>> subset_results = results[1:4] # Get a slice of results
|
287
|
+
"""
|
126
288
|
return self._apply("__getitem__", idx)
|
127
289
|
|
128
290
|
def __len__(self):
|
129
|
-
"""
|
291
|
+
"""
|
292
|
+
Return the number of detections in the Results object.
|
293
|
+
|
294
|
+
Returns:
|
295
|
+
(int): The number of detections, determined by the length of the first non-empty attribute
|
296
|
+
(boxes, masks, probs, keypoints, or obb).
|
297
|
+
|
298
|
+
Examples:
|
299
|
+
>>> results = Results(orig_img, path, names, boxes=torch.rand(5, 4))
|
300
|
+
>>> len(results)
|
301
|
+
5
|
302
|
+
"""
|
130
303
|
for k in self._keys:
|
131
304
|
v = getattr(self, k)
|
132
305
|
if v is not None:
|
133
306
|
return len(v)
|
134
307
|
|
135
308
|
def update(self, boxes=None, masks=None, probs=None, obb=None):
|
136
|
-
"""
|
309
|
+
"""
|
310
|
+
Updates the Results object with new detection data.
|
311
|
+
|
312
|
+
This method allows updating the boxes, masks, probabilities, and oriented bounding boxes (OBB) of the
|
313
|
+
Results object. It ensures that boxes are clipped to the original image shape.
|
314
|
+
|
315
|
+
Args:
|
316
|
+
boxes (torch.Tensor | None): A tensor of shape (N, 6) containing bounding box coordinates and
|
317
|
+
confidence scores. The format is (x1, y1, x2, y2, conf, class).
|
318
|
+
masks (torch.Tensor | None): A tensor of shape (N, H, W) containing segmentation masks.
|
319
|
+
probs (torch.Tensor | None): A tensor of shape (num_classes,) containing class probabilities.
|
320
|
+
obb (torch.Tensor | None): A tensor of shape (N, 5) containing oriented bounding box coordinates.
|
321
|
+
|
322
|
+
Examples:
|
323
|
+
>>> results = model("image.jpg")
|
324
|
+
>>> new_boxes = torch.tensor([[100, 100, 200, 200, 0.9, 0]])
|
325
|
+
>>> results[0].update(boxes=new_boxes)
|
326
|
+
"""
|
137
327
|
if boxes is not None:
|
138
328
|
self.boxes = Boxes(ops.clip_boxes(boxes, self.orig_shape), self.orig_shape)
|
139
329
|
if masks is not None:
|
@@ -145,16 +335,23 @@ class Results(SimpleClass):
|
|
145
335
|
|
146
336
|
def _apply(self, fn, *args, **kwargs):
|
147
337
|
"""
|
148
|
-
Applies a function to all non-empty attributes and returns a new Results object with modified attributes.
|
149
|
-
|
338
|
+
Applies a function to all non-empty attributes and returns a new Results object with modified attributes.
|
339
|
+
|
340
|
+
This method is internally called by methods like .to(), .cuda(), .cpu(), etc.
|
150
341
|
|
151
342
|
Args:
|
152
343
|
fn (str): The name of the function to apply.
|
153
|
-
*args: Variable length argument list to pass to the function.
|
154
|
-
**kwargs: Arbitrary keyword arguments to pass to the function.
|
344
|
+
*args (Any): Variable length argument list to pass to the function.
|
345
|
+
**kwargs (Any): Arbitrary keyword arguments to pass to the function.
|
155
346
|
|
156
347
|
Returns:
|
157
|
-
Results: A new Results object with attributes modified by the applied function.
|
348
|
+
(Results): A new Results object with attributes modified by the applied function.
|
349
|
+
|
350
|
+
Examples:
|
351
|
+
>>> results = model("path/to/image.jpg")
|
352
|
+
>>> for result in results:
|
353
|
+
... result_cuda = result.cuda()
|
354
|
+
... result_cpu = result.cpu()
|
158
355
|
"""
|
159
356
|
r = self.new()
|
160
357
|
for k in self._keys:
|
@@ -164,24 +361,87 @@ class Results(SimpleClass):
|
|
164
361
|
return r
|
165
362
|
|
166
363
|
def cpu(self):
|
167
|
-
"""
|
364
|
+
"""
|
365
|
+
Returns a copy of the Results object with all its tensors moved to CPU memory.
|
366
|
+
|
367
|
+
This method creates a new Results object with all tensor attributes (boxes, masks, probs, keypoints, obb)
|
368
|
+
transferred to CPU memory. It's useful for moving data from GPU to CPU for further processing or saving.
|
369
|
+
|
370
|
+
Returns:
|
371
|
+
(Results): A new Results object with all tensor attributes on CPU memory.
|
372
|
+
|
373
|
+
Examples:
|
374
|
+
>>> results = model("path/to/image.jpg") # Perform inference
|
375
|
+
>>> cpu_result = results[0].cpu() # Move the first result to CPU
|
376
|
+
>>> print(cpu_result.boxes.device) # Output: cpu
|
377
|
+
"""
|
168
378
|
return self._apply("cpu")
|
169
379
|
|
170
380
|
def numpy(self):
|
171
|
-
"""
|
381
|
+
"""
|
382
|
+
Converts all tensors in the Results object to numpy arrays.
|
383
|
+
|
384
|
+
Returns:
|
385
|
+
(Results): A new Results object with all tensors converted to numpy arrays.
|
386
|
+
|
387
|
+
Examples:
|
388
|
+
>>> results = model("path/to/image.jpg")
|
389
|
+
>>> numpy_result = results[0].numpy()
|
390
|
+
>>> type(numpy_result.boxes.data)
|
391
|
+
<class 'numpy.ndarray'>
|
392
|
+
|
393
|
+
Notes:
|
394
|
+
This method creates a new Results object, leaving the original unchanged. It's useful for
|
395
|
+
interoperability with numpy-based libraries or when CPU-based operations are required.
|
396
|
+
"""
|
172
397
|
return self._apply("numpy")
|
173
398
|
|
174
399
|
def cuda(self):
|
175
|
-
"""
|
400
|
+
"""
|
401
|
+
Moves all tensors in the Results object to GPU memory.
|
402
|
+
|
403
|
+
Returns:
|
404
|
+
(Results): A new Results object with all tensors moved to CUDA device.
|
405
|
+
|
406
|
+
Examples:
|
407
|
+
>>> results = model("path/to/image.jpg")
|
408
|
+
>>> cuda_results = results[0].cuda() # Move first result to GPU
|
409
|
+
>>> for result in results:
|
410
|
+
... result_cuda = result.cuda() # Move each result to GPU
|
411
|
+
"""
|
176
412
|
return self._apply("cuda")
|
177
413
|
|
178
414
|
def to(self, *args, **kwargs):
|
179
|
-
"""
|
415
|
+
"""
|
416
|
+
Moves all tensors in the Results object to the specified device and dtype.
|
417
|
+
|
418
|
+
Args:
|
419
|
+
*args (Any): Variable length argument list to be passed to torch.Tensor.to().
|
420
|
+
**kwargs (Any): Arbitrary keyword arguments to be passed to torch.Tensor.to().
|
421
|
+
|
422
|
+
Returns:
|
423
|
+
(Results): A new Results object with all tensors moved to the specified device and dtype.
|
424
|
+
|
425
|
+
Examples:
|
426
|
+
>>> results = model("path/to/image.jpg")
|
427
|
+
>>> result_cuda = results[0].to("cuda") # Move first result to GPU
|
428
|
+
>>> result_cpu = results[0].to("cpu") # Move first result to CPU
|
429
|
+
>>> result_half = results[0].to(dtype=torch.float16) # Convert first result to half precision
|
430
|
+
"""
|
180
431
|
return self._apply("to", *args, **kwargs)
|
181
432
|
|
182
433
|
def new(self):
|
183
|
-
"""
|
184
|
-
|
434
|
+
"""
|
435
|
+
Creates a new Results object with the same image, path, names, and speed attributes.
|
436
|
+
|
437
|
+
Returns:
|
438
|
+
(Results): A new Results object with copied attributes from the original instance.
|
439
|
+
|
440
|
+
Examples:
|
441
|
+
>>> results = model("path/to/image.jpg")
|
442
|
+
>>> new_result = results[0].new()
|
443
|
+
"""
|
444
|
+
return Results(orig_img=self.orig_img, path=self.path, names=self.names, speed=self.speed)
|
185
445
|
|
186
446
|
def plot(
|
187
447
|
self,
|
@@ -201,45 +461,40 @@ class Results(SimpleClass):
|
|
201
461
|
show=False,
|
202
462
|
save=False,
|
203
463
|
filename=None,
|
464
|
+
color_mode="class",
|
204
465
|
):
|
205
466
|
"""
|
206
|
-
Plots
|
467
|
+
Plots detection results on an input RGB image.
|
207
468
|
|
208
469
|
Args:
|
209
|
-
conf (bool): Whether to plot
|
210
|
-
line_width (float
|
211
|
-
font_size (float
|
212
|
-
font (str):
|
470
|
+
conf (bool): Whether to plot detection confidence scores.
|
471
|
+
line_width (float | None): Line width of bounding boxes. If None, scaled to image size.
|
472
|
+
font_size (float | None): Font size for text. If None, scaled to image size.
|
473
|
+
font (str): Font to use for text.
|
213
474
|
pil (bool): Whether to return the image as a PIL Image.
|
214
|
-
img (
|
215
|
-
im_gpu (torch.Tensor): Normalized image
|
216
|
-
kpt_radius (int
|
475
|
+
img (np.ndarray | None): Image to plot on. If None, uses original image.
|
476
|
+
im_gpu (torch.Tensor | None): Normalized image on GPU for faster mask plotting.
|
477
|
+
kpt_radius (int): Radius of drawn keypoints.
|
217
478
|
kpt_line (bool): Whether to draw lines connecting keypoints.
|
218
|
-
labels (bool): Whether to plot
|
219
|
-
boxes (bool): Whether to plot
|
220
|
-
masks (bool): Whether to plot
|
221
|
-
probs (bool): Whether to plot classification
|
222
|
-
show (bool): Whether to display the annotated image
|
223
|
-
save (bool): Whether to save the annotated image
|
224
|
-
filename (str): Filename to save image
|
225
|
-
|
226
|
-
|
227
|
-
|
228
|
-
|
229
|
-
|
230
|
-
|
231
|
-
|
232
|
-
|
233
|
-
|
234
|
-
|
235
|
-
results = model('bus.jpg') # results list
|
236
|
-
for r in results:
|
237
|
-
im_array = r.plot() # plot a BGR numpy array of predictions
|
238
|
-
im = Image.fromarray(im_array[..., ::-1]) # RGB PIL image
|
239
|
-
im.show() # show image
|
240
|
-
im.save('results.jpg') # save image
|
241
|
-
```
|
479
|
+
labels (bool): Whether to plot labels of bounding boxes.
|
480
|
+
boxes (bool): Whether to plot bounding boxes.
|
481
|
+
masks (bool): Whether to plot masks.
|
482
|
+
probs (bool): Whether to plot classification probabilities.
|
483
|
+
show (bool): Whether to display the annotated image.
|
484
|
+
save (bool): Whether to save the annotated image.
|
485
|
+
filename (str | None): Filename to save image if save is True.
|
486
|
+
color_mode (bool): Specify the color mode, e.g., 'instance' or 'class'. Default to 'class'.
|
487
|
+
|
488
|
+
Returns:
|
489
|
+
(np.ndarray): Annotated image as a numpy array.
|
490
|
+
|
491
|
+
Examples:
|
492
|
+
>>> results = model("image.jpg")
|
493
|
+
>>> for result in results:
|
494
|
+
... im = result.plot()
|
495
|
+
... im.show()
|
242
496
|
"""
|
497
|
+
assert color_mode in {"instance", "class"}, f"Expected color_mode='instance' or 'class', not {color_mode}."
|
243
498
|
if img is None and isinstance(self.orig_img, torch.Tensor):
|
244
499
|
img = (self.orig_img[0].detach().permute(1, 2, 0).contiguous() * 255).to(torch.uint8).cpu().numpy()
|
245
500
|
|
@@ -268,17 +523,37 @@ class Results(SimpleClass):
|
|
268
523
|
.contiguous()
|
269
524
|
/ 255
|
270
525
|
)
|
271
|
-
idx =
|
526
|
+
idx = (
|
527
|
+
pred_boxes.id
|
528
|
+
if pred_boxes.id is not None and color_mode == "instance"
|
529
|
+
else pred_boxes.cls
|
530
|
+
if pred_boxes and color_mode == "class"
|
531
|
+
else reversed(range(len(pred_masks)))
|
532
|
+
)
|
272
533
|
annotator.masks(pred_masks.data, colors=[colors(x, True) for x in idx], im_gpu=im_gpu)
|
273
534
|
|
274
535
|
# Plot Detect results
|
275
536
|
if pred_boxes is not None and show_boxes:
|
276
|
-
for d in reversed(pred_boxes):
|
277
|
-
c,
|
537
|
+
for i, d in enumerate(reversed(pred_boxes)):
|
538
|
+
c, d_conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())
|
278
539
|
name = ("" if id is None else f"id:{id} ") + names[c]
|
279
|
-
label = (f"{name} {
|
540
|
+
label = (f"{name} {d_conf:.2f}" if conf else name) if labels else None
|
280
541
|
box = d.xyxyxyxy.reshape(-1, 4, 2).squeeze() if is_obb else d.xyxy.squeeze()
|
281
|
-
annotator.box_label(
|
542
|
+
annotator.box_label(
|
543
|
+
box,
|
544
|
+
label,
|
545
|
+
color=colors(
|
546
|
+
c
|
547
|
+
if color_mode == "class"
|
548
|
+
else id
|
549
|
+
if id is not None
|
550
|
+
else i
|
551
|
+
if color_mode == "instance"
|
552
|
+
else None,
|
553
|
+
True,
|
554
|
+
),
|
555
|
+
rotated=is_obb,
|
556
|
+
)
|
282
557
|
|
283
558
|
# Plot Classify results
|
284
559
|
if pred_probs is not None and show_probs:
|
@@ -288,8 +563,14 @@ class Results(SimpleClass):
|
|
288
563
|
|
289
564
|
# Plot Pose results
|
290
565
|
if self.keypoints is not None:
|
291
|
-
for k in reversed(self.keypoints.data):
|
292
|
-
annotator.kpts(
|
566
|
+
for i, k in enumerate(reversed(self.keypoints.data)):
|
567
|
+
annotator.kpts(
|
568
|
+
k,
|
569
|
+
self.orig_shape,
|
570
|
+
radius=kpt_radius,
|
571
|
+
kpt_line=kpt_line,
|
572
|
+
kpt_color=colors(i, True) if color_mode == "instance" else None,
|
573
|
+
)
|
293
574
|
|
294
575
|
# Show results
|
295
576
|
if show:
|
@@ -302,26 +583,80 @@ class Results(SimpleClass):
|
|
302
583
|
return annotator.result()
|
303
584
|
|
304
585
|
def show(self, *args, **kwargs):
|
305
|
-
"""
|
586
|
+
"""
|
587
|
+
Display the image with annotated inference results.
|
588
|
+
|
589
|
+
This method plots the detection results on the original image and displays it. It's a convenient way to
|
590
|
+
visualize the model's predictions directly.
|
591
|
+
|
592
|
+
Args:
|
593
|
+
*args (Any): Variable length argument list to be passed to the `plot()` method.
|
594
|
+
**kwargs (Any): Arbitrary keyword arguments to be passed to the `plot()` method.
|
595
|
+
|
596
|
+
Examples:
|
597
|
+
>>> results = model("path/to/image.jpg")
|
598
|
+
>>> results[0].show() # Display the first result
|
599
|
+
>>> for result in results:
|
600
|
+
... result.show() # Display all results
|
601
|
+
"""
|
306
602
|
self.plot(show=True, *args, **kwargs)
|
307
603
|
|
308
604
|
def save(self, filename=None, *args, **kwargs):
|
309
|
-
"""
|
605
|
+
"""
|
606
|
+
Saves annotated inference results image to file.
|
607
|
+
|
608
|
+
This method plots the detection results on the original image and saves the annotated image to a file. It
|
609
|
+
utilizes the `plot` method to generate the annotated image and then saves it to the specified filename.
|
610
|
+
|
611
|
+
Args:
|
612
|
+
filename (str | Path | None): The filename to save the annotated image. If None, a default filename
|
613
|
+
is generated based on the original image path.
|
614
|
+
*args (Any): Variable length argument list to be passed to the `plot` method.
|
615
|
+
**kwargs (Any): Arbitrary keyword arguments to be passed to the `plot` method.
|
616
|
+
|
617
|
+
Examples:
|
618
|
+
>>> results = model("path/to/image.jpg")
|
619
|
+
>>> for result in results:
|
620
|
+
... result.save("annotated_image.jpg")
|
621
|
+
>>> # Or with custom plot arguments
|
622
|
+
>>> for result in results:
|
623
|
+
... result.save("annotated_image.jpg", conf=False, line_width=2)
|
624
|
+
"""
|
310
625
|
if not filename:
|
311
626
|
filename = f"results_{Path(self.path).name}"
|
312
627
|
self.plot(save=True, filename=filename, *args, **kwargs)
|
313
628
|
return filename
|
314
629
|
|
315
630
|
def verbose(self):
|
316
|
-
"""
|
631
|
+
"""
|
632
|
+
Returns a log string for each task in the results, detailing detection and classification outcomes.
|
633
|
+
|
634
|
+
This method generates a human-readable string summarizing the detection and classification results. It includes
|
635
|
+
the number of detections for each class and the top probabilities for classification tasks.
|
636
|
+
|
637
|
+
Returns:
|
638
|
+
(str): A formatted string containing a summary of the results. For detection tasks, it includes the
|
639
|
+
number of detections per class. For classification tasks, it includes the top 5 class probabilities.
|
640
|
+
|
641
|
+
Examples:
|
642
|
+
>>> results = model("path/to/image.jpg")
|
643
|
+
>>> for result in results:
|
644
|
+
... print(result.verbose())
|
645
|
+
2 persons, 1 car, 3 traffic lights,
|
646
|
+
dog 0.92, cat 0.78, horse 0.64,
|
647
|
+
|
648
|
+
Notes:
|
649
|
+
- If there are no detections, the method returns "(no detections), " for detection tasks.
|
650
|
+
- For classification tasks, it returns the top 5 class probabilities and their corresponding class names.
|
651
|
+
- The returned string is comma-separated and ends with a comma and a space.
|
652
|
+
"""
|
317
653
|
log_string = ""
|
318
654
|
probs = self.probs
|
319
|
-
boxes = self.boxes
|
320
655
|
if len(self) == 0:
|
321
656
|
return log_string if probs is not None else f"{log_string}(no detections), "
|
322
657
|
if probs is not None:
|
323
658
|
log_string += f"{', '.join(f'{self.names[j]} {probs.data[j]:.2f}' for j in probs.top5)}, "
|
324
|
-
if boxes:
|
659
|
+
if boxes := self.boxes:
|
325
660
|
for c in boxes.cls.unique():
|
326
661
|
n = (boxes.cls == c).sum() # detections per class
|
327
662
|
log_string += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "
|
@@ -329,11 +664,30 @@ class Results(SimpleClass):
|
|
329
664
|
|
330
665
|
def save_txt(self, txt_file, save_conf=False):
|
331
666
|
"""
|
332
|
-
Save
|
667
|
+
Save detection results to a text file.
|
333
668
|
|
334
669
|
Args:
|
335
|
-
txt_file (str):
|
336
|
-
save_conf (bool):
|
670
|
+
txt_file (str | Path): Path to the output text file.
|
671
|
+
save_conf (bool): Whether to include confidence scores in the output.
|
672
|
+
|
673
|
+
Returns:
|
674
|
+
(str): Path to the saved text file.
|
675
|
+
|
676
|
+
Examples:
|
677
|
+
>>> from ultralytics import YOLO
|
678
|
+
>>> model = YOLO("yolo11n.pt")
|
679
|
+
>>> results = model("path/to/image.jpg")
|
680
|
+
>>> for result in results:
|
681
|
+
... result.save_txt("output.txt")
|
682
|
+
|
683
|
+
Notes:
|
684
|
+
- The file will contain one line per detection or classification with the following structure:
|
685
|
+
- For detections: `class confidence x_center y_center width height`
|
686
|
+
- For classifications: `confidence class_name`
|
687
|
+
- For masks and keypoints, the specific formats will vary accordingly.
|
688
|
+
- The function will create the output directory if it does not exist.
|
689
|
+
- If save_conf is False, the confidence scores will be excluded from the output.
|
690
|
+
- Existing contents of the file will not be overwritten; new results will be appended.
|
337
691
|
"""
|
338
692
|
is_obb = self.obb is not None
|
339
693
|
boxes = self.obb if is_obb else self.boxes
|
@@ -365,11 +719,25 @@ class Results(SimpleClass):
|
|
365
719
|
|
366
720
|
def save_crop(self, save_dir, file_name=Path("im.jpg")):
|
367
721
|
"""
|
368
|
-
|
722
|
+
Saves cropped detection images to specified directory.
|
723
|
+
|
724
|
+
This method saves cropped images of detected objects to a specified directory. Each crop is saved in a
|
725
|
+
subdirectory named after the object's class, with the filename based on the input file_name.
|
369
726
|
|
370
727
|
Args:
|
371
|
-
save_dir (str |
|
372
|
-
file_name (str |
|
728
|
+
save_dir (str | Path): Directory path where cropped images will be saved.
|
729
|
+
file_name (str | Path): Base filename for the saved cropped images. Default is Path("im.jpg").
|
730
|
+
|
731
|
+
Notes:
|
732
|
+
- This method does not support Classify or Oriented Bounding Box (OBB) tasks.
|
733
|
+
- Crops are saved as 'save_dir/class_name/file_name.jpg'.
|
734
|
+
- The method will create necessary subdirectories if they don't exist.
|
735
|
+
- Original image is copied before cropping to avoid modifying the original.
|
736
|
+
|
737
|
+
Examples:
|
738
|
+
>>> results = model("path/to/image.jpg")
|
739
|
+
>>> for result in results:
|
740
|
+
... result.save_crop(save_dir="path/to/crops", file_name="detection")
|
373
741
|
"""
|
374
742
|
if self.probs is not None:
|
375
743
|
LOGGER.warning("WARNING ⚠️ Classify task do not support `save_crop`.")
|
@@ -381,117 +749,379 @@ class Results(SimpleClass):
|
|
381
749
|
save_one_box(
|
382
750
|
d.xyxy,
|
383
751
|
self.orig_img.copy(),
|
384
|
-
file=Path(save_dir) / self.names[int(d.cls)] /
|
752
|
+
file=Path(save_dir) / self.names[int(d.cls)] / Path(file_name).with_suffix(".jpg"),
|
385
753
|
BGR=True,
|
386
754
|
)
|
387
755
|
|
388
|
-
def
|
389
|
-
"""
|
390
|
-
|
391
|
-
LOGGER.warning("Warning: Classify task do not support `tojson` yet.")
|
392
|
-
return
|
756
|
+
def summary(self, normalize=False, decimals=5):
|
757
|
+
"""
|
758
|
+
Converts inference results to a summarized dictionary with optional normalization for box coordinates.
|
393
759
|
|
394
|
-
|
760
|
+
This method creates a list of detection dictionaries, each containing information about a single
|
761
|
+
detection or classification result. For classification tasks, it returns the top class and its
|
762
|
+
confidence. For detection tasks, it includes class information, bounding box coordinates, and
|
763
|
+
optionally mask segments and keypoints.
|
395
764
|
|
765
|
+
Args:
|
766
|
+
normalize (bool): Whether to normalize bounding box coordinates by image dimensions. Defaults to False.
|
767
|
+
decimals (int): Number of decimal places to round the output values to. Defaults to 5.
|
768
|
+
|
769
|
+
Returns:
|
770
|
+
(List[Dict]): A list of dictionaries, each containing summarized information for a single
|
771
|
+
detection or classification result. The structure of each dictionary varies based on the
|
772
|
+
task type (classification or detection) and available information (boxes, masks, keypoints).
|
773
|
+
|
774
|
+
Examples:
|
775
|
+
>>> results = model("image.jpg")
|
776
|
+
>>> summary = results[0].summary()
|
777
|
+
>>> print(summary)
|
778
|
+
"""
|
396
779
|
# Create list of detection dictionaries
|
397
780
|
results = []
|
398
|
-
|
781
|
+
if self.probs is not None:
|
782
|
+
class_id = self.probs.top1
|
783
|
+
results.append(
|
784
|
+
{
|
785
|
+
"name": self.names[class_id],
|
786
|
+
"class": class_id,
|
787
|
+
"confidence": round(self.probs.top1conf.item(), decimals),
|
788
|
+
}
|
789
|
+
)
|
790
|
+
return results
|
791
|
+
|
792
|
+
is_obb = self.obb is not None
|
793
|
+
data = self.obb if is_obb else self.boxes
|
399
794
|
h, w = self.orig_shape if normalize else (1, 1)
|
400
795
|
for i, row in enumerate(data): # xyxy, track_id if tracking, conf, class_id
|
401
|
-
|
402
|
-
|
403
|
-
|
404
|
-
|
405
|
-
|
406
|
-
|
407
|
-
|
796
|
+
class_id, conf = int(row.cls), round(row.conf.item(), decimals)
|
797
|
+
box = (row.xyxyxyxy if is_obb else row.xyxy).squeeze().reshape(-1, 2).tolist()
|
798
|
+
xy = {}
|
799
|
+
for j, b in enumerate(box):
|
800
|
+
xy[f"x{j + 1}"] = round(b[0] / w, decimals)
|
801
|
+
xy[f"y{j + 1}"] = round(b[1] / h, decimals)
|
802
|
+
result = {"name": self.names[class_id], "class": class_id, "confidence": conf, "box": xy}
|
803
|
+
if data.is_track:
|
804
|
+
result["track_id"] = int(row.id.item()) # track ID
|
408
805
|
if self.masks:
|
409
|
-
|
410
|
-
|
806
|
+
result["segments"] = {
|
807
|
+
"x": (self.masks.xy[i][:, 0] / w).round(decimals).tolist(),
|
808
|
+
"y": (self.masks.xy[i][:, 1] / h).round(decimals).tolist(),
|
809
|
+
}
|
411
810
|
if self.keypoints is not None:
|
412
811
|
x, y, visible = self.keypoints[i].data[0].cpu().unbind(dim=1) # torch Tensor
|
413
|
-
result["keypoints"] = {
|
812
|
+
result["keypoints"] = {
|
813
|
+
"x": (x / w).numpy().round(decimals).tolist(), # decimals named argument required
|
814
|
+
"y": (y / h).numpy().round(decimals).tolist(),
|
815
|
+
"visible": visible.numpy().round(decimals).tolist(),
|
816
|
+
}
|
414
817
|
results.append(result)
|
415
818
|
|
416
|
-
|
417
|
-
|
819
|
+
return results
|
820
|
+
|
821
|
+
def to_df(self, normalize=False, decimals=5):
|
822
|
+
"""
|
823
|
+
Converts detection results to a Pandas Dataframe.
|
824
|
+
|
825
|
+
This method converts the detection results into Pandas Dataframe format. It includes information
|
826
|
+
about detected objects such as bounding boxes, class names, confidence scores, and optionally
|
827
|
+
segmentation masks and keypoints.
|
828
|
+
|
829
|
+
Args:
|
830
|
+
normalize (bool): Whether to normalize the bounding box coordinates by the image dimensions.
|
831
|
+
If True, coordinates will be returned as float values between 0 and 1. Defaults to False.
|
832
|
+
decimals (int): Number of decimal places to round the output values to. Defaults to 5.
|
833
|
+
|
834
|
+
Returns:
|
835
|
+
(DataFrame): A Pandas Dataframe containing all the information in results in an organized way.
|
836
|
+
|
837
|
+
Examples:
|
838
|
+
>>> results = model("path/to/image.jpg")
|
839
|
+
>>> df_result = results[0].to_df()
|
840
|
+
>>> print(df_result)
|
841
|
+
"""
|
842
|
+
import pandas as pd # scope for faster 'import ultralytics'
|
843
|
+
|
844
|
+
return pd.DataFrame(self.summary(normalize=normalize, decimals=decimals))
|
845
|
+
|
846
|
+
def to_csv(self, normalize=False, decimals=5, *args, **kwargs):
|
847
|
+
"""
|
848
|
+
Converts detection results to a CSV format.
|
849
|
+
|
850
|
+
This method serializes the detection results into a CSV format. It includes information
|
851
|
+
about detected objects such as bounding boxes, class names, confidence scores, and optionally
|
852
|
+
segmentation masks and keypoints.
|
853
|
+
|
854
|
+
Args:
|
855
|
+
normalize (bool): Whether to normalize the bounding box coordinates by the image dimensions.
|
856
|
+
If True, coordinates will be returned as float values between 0 and 1. Defaults to False.
|
857
|
+
decimals (int): Number of decimal places to round the output values to. Defaults to 5.
|
858
|
+
*args (Any): Variable length argument list to be passed to pandas.DataFrame.to_csv().
|
859
|
+
**kwargs (Any): Arbitrary keyword arguments to be passed to pandas.DataFrame.to_csv().
|
860
|
+
|
861
|
+
|
862
|
+
Returns:
|
863
|
+
(str): CSV containing all the information in results in an organized way.
|
864
|
+
|
865
|
+
Examples:
|
866
|
+
>>> results = model("path/to/image.jpg")
|
867
|
+
>>> csv_result = results[0].to_csv()
|
868
|
+
>>> print(csv_result)
|
869
|
+
"""
|
870
|
+
return self.to_df(normalize=normalize, decimals=decimals).to_csv(*args, **kwargs)
|
871
|
+
|
872
|
+
def to_xml(self, normalize=False, decimals=5, *args, **kwargs):
|
873
|
+
"""
|
874
|
+
Converts detection results to XML format.
|
875
|
+
|
876
|
+
This method serializes the detection results into an XML format. It includes information
|
877
|
+
about detected objects such as bounding boxes, class names, confidence scores, and optionally
|
878
|
+
segmentation masks and keypoints.
|
879
|
+
|
880
|
+
Args:
|
881
|
+
normalize (bool): Whether to normalize the bounding box coordinates by the image dimensions.
|
882
|
+
If True, coordinates will be returned as float values between 0 and 1. Defaults to False.
|
883
|
+
decimals (int): Number of decimal places to round the output values to. Defaults to 5.
|
884
|
+
*args (Any): Variable length argument list to be passed to pandas.DataFrame.to_xml().
|
885
|
+
**kwargs (Any): Arbitrary keyword arguments to be passed to pandas.DataFrame.to_xml().
|
886
|
+
|
887
|
+
Returns:
|
888
|
+
(str): An XML string containing all the information in results in an organized way.
|
889
|
+
|
890
|
+
Examples:
|
891
|
+
>>> results = model("path/to/image.jpg")
|
892
|
+
>>> xml_result = results[0].to_xml()
|
893
|
+
>>> print(xml_result)
|
894
|
+
"""
|
895
|
+
check_requirements("lxml")
|
896
|
+
df = self.to_df(normalize=normalize, decimals=decimals)
|
897
|
+
return '<?xml version="1.0" encoding="utf-8"?>\n<root></root>' if df.empty else df.to_xml(*args, **kwargs)
|
898
|
+
|
899
|
+
def tojson(self, normalize=False, decimals=5):
|
900
|
+
"""Deprecated version of to_json()."""
|
901
|
+
LOGGER.warning("WARNING ⚠️ 'result.tojson()' is deprecated, replace with 'result.to_json()'.")
|
902
|
+
return self.to_json(normalize, decimals)
|
903
|
+
|
904
|
+
def to_json(self, normalize=False, decimals=5):
|
905
|
+
"""
|
906
|
+
Converts detection results to JSON format.
|
907
|
+
|
908
|
+
This method serializes the detection results into a JSON-compatible format. It includes information
|
909
|
+
about detected objects such as bounding boxes, class names, confidence scores, and optionally
|
910
|
+
segmentation masks and keypoints.
|
911
|
+
|
912
|
+
Args:
|
913
|
+
normalize (bool): Whether to normalize the bounding box coordinates by the image dimensions.
|
914
|
+
If True, coordinates will be returned as float values between 0 and 1. Defaults to False.
|
915
|
+
decimals (int): Number of decimal places to round the output values to. Defaults to 5.
|
916
|
+
|
917
|
+
Returns:
|
918
|
+
(str): A JSON string containing the serialized detection results.
|
919
|
+
|
920
|
+
Examples:
|
921
|
+
>>> results = model("path/to/image.jpg")
|
922
|
+
>>> json_result = results[0].to_json()
|
923
|
+
>>> print(json_result)
|
924
|
+
|
925
|
+
Notes:
|
926
|
+
- For classification tasks, the JSON will contain class probabilities instead of bounding boxes.
|
927
|
+
- For object detection tasks, the JSON will include bounding box coordinates, class names, and
|
928
|
+
confidence scores.
|
929
|
+
- If available, segmentation masks and keypoints will also be included in the JSON output.
|
930
|
+
- The method uses the `summary` method internally to generate the data structure before
|
931
|
+
converting it to JSON.
|
932
|
+
"""
|
933
|
+
import json
|
934
|
+
|
935
|
+
return json.dumps(self.summary(normalize=normalize, decimals=decimals), indent=2)
|
418
936
|
|
419
937
|
|
420
938
|
class Boxes(BaseTensor):
|
421
939
|
"""
|
422
|
-
|
423
|
-
|
424
|
-
|
940
|
+
A class for managing and manipulating detection boxes.
|
941
|
+
|
942
|
+
This class provides functionality for handling detection boxes, including their coordinates, confidence scores,
|
943
|
+
class labels, and optional tracking IDs. It supports various box formats and offers methods for easy manipulation
|
944
|
+
and conversion between different coordinate systems.
|
425
945
|
|
426
946
|
Attributes:
|
427
|
-
data (torch.Tensor): The raw tensor containing detection boxes and
|
428
|
-
orig_shape (
|
947
|
+
data (torch.Tensor | numpy.ndarray): The raw tensor containing detection boxes and associated data.
|
948
|
+
orig_shape (Tuple[int, int]): The original image dimensions (height, width).
|
429
949
|
is_track (bool): Indicates whether tracking IDs are included in the box data.
|
430
|
-
|
431
|
-
Properties:
|
432
950
|
xyxy (torch.Tensor | numpy.ndarray): Boxes in [x1, y1, x2, y2] format.
|
433
951
|
conf (torch.Tensor | numpy.ndarray): Confidence scores for each box.
|
434
952
|
cls (torch.Tensor | numpy.ndarray): Class labels for each box.
|
435
|
-
id (torch.Tensor | numpy.ndarray
|
436
|
-
xywh (torch.Tensor | numpy.ndarray): Boxes in [x, y, width, height] format
|
437
|
-
xyxyn (torch.Tensor | numpy.ndarray): Normalized [x1, y1, x2, y2] boxes
|
438
|
-
xywhn (torch.Tensor | numpy.ndarray): Normalized [x, y, width, height] boxes
|
953
|
+
id (torch.Tensor | numpy.ndarray): Tracking IDs for each box (if available).
|
954
|
+
xywh (torch.Tensor | numpy.ndarray): Boxes in [x, y, width, height] format.
|
955
|
+
xyxyn (torch.Tensor | numpy.ndarray): Normalized [x1, y1, x2, y2] boxes relative to orig_shape.
|
956
|
+
xywhn (torch.Tensor | numpy.ndarray): Normalized [x, y, width, height] boxes relative to orig_shape.
|
439
957
|
|
440
958
|
Methods:
|
441
|
-
cpu():
|
442
|
-
numpy():
|
443
|
-
cuda():
|
444
|
-
to(
|
959
|
+
cpu(): Returns a copy of the object with all tensors on CPU memory.
|
960
|
+
numpy(): Returns a copy of the object with all tensors as numpy arrays.
|
961
|
+
cuda(): Returns a copy of the object with all tensors on GPU memory.
|
962
|
+
to(*args, **kwargs): Returns a copy of the object with tensors on specified device and dtype.
|
963
|
+
|
964
|
+
Examples:
|
965
|
+
>>> import torch
|
966
|
+
>>> boxes_data = torch.tensor([[100, 50, 150, 100, 0.9, 0], [200, 150, 300, 250, 0.8, 1]])
|
967
|
+
>>> orig_shape = (480, 640) # height, width
|
968
|
+
>>> boxes = Boxes(boxes_data, orig_shape)
|
969
|
+
>>> print(boxes.xyxy)
|
970
|
+
>>> print(boxes.conf)
|
971
|
+
>>> print(boxes.cls)
|
972
|
+
>>> print(boxes.xywhn)
|
445
973
|
"""
|
446
974
|
|
447
975
|
def __init__(self, boxes, orig_shape) -> None:
|
448
976
|
"""
|
449
|
-
Initialize the Boxes class.
|
977
|
+
Initialize the Boxes class with detection box data and the original image shape.
|
978
|
+
|
979
|
+
This class manages detection boxes, providing easy access and manipulation of box coordinates,
|
980
|
+
confidence scores, class identifiers, and optional tracking IDs. It supports multiple formats
|
981
|
+
for box coordinates, including both absolute and normalized forms.
|
450
982
|
|
451
983
|
Args:
|
452
|
-
boxes (torch.Tensor |
|
453
|
-
|
454
|
-
|
455
|
-
orig_shape (
|
984
|
+
boxes (torch.Tensor | np.ndarray): A tensor or numpy array with detection boxes of shape
|
985
|
+
(num_boxes, 6) or (num_boxes, 7). Columns should contain
|
986
|
+
[x1, y1, x2, y2, confidence, class, (optional) track_id].
|
987
|
+
orig_shape (Tuple[int, int]): The original image shape as (height, width). Used for normalization.
|
988
|
+
|
989
|
+
Attributes:
|
990
|
+
data (torch.Tensor): The raw tensor containing detection boxes and their associated data.
|
991
|
+
orig_shape (Tuple[int, int]): The original image size, used for normalization.
|
992
|
+
is_track (bool): Indicates whether tracking IDs are included in the box data.
|
993
|
+
|
994
|
+
Examples:
|
995
|
+
>>> import torch
|
996
|
+
>>> boxes = torch.tensor([[100, 50, 150, 100, 0.9, 0]])
|
997
|
+
>>> orig_shape = (480, 640)
|
998
|
+
>>> detection_boxes = Boxes(boxes, orig_shape)
|
999
|
+
>>> print(detection_boxes.xyxy)
|
1000
|
+
tensor([[100., 50., 150., 100.]])
|
456
1001
|
"""
|
457
1002
|
if boxes.ndim == 1:
|
458
1003
|
boxes = boxes[None, :]
|
459
1004
|
n = boxes.shape[-1]
|
460
|
-
assert n in
|
1005
|
+
assert n in {6, 7}, f"expected 6 or 7 values but got {n}" # xyxy, track_id, conf, cls
|
461
1006
|
super().__init__(boxes, orig_shape)
|
462
1007
|
self.is_track = n == 7
|
463
1008
|
self.orig_shape = orig_shape
|
464
1009
|
|
465
1010
|
@property
|
466
1011
|
def xyxy(self):
|
467
|
-
"""
|
1012
|
+
"""
|
1013
|
+
Returns bounding boxes in [x1, y1, x2, y2] format.
|
1014
|
+
|
1015
|
+
Returns:
|
1016
|
+
(torch.Tensor | numpy.ndarray): A tensor or numpy array of shape (n, 4) containing bounding box
|
1017
|
+
coordinates in [x1, y1, x2, y2] format, where n is the number of boxes.
|
1018
|
+
|
1019
|
+
Examples:
|
1020
|
+
>>> results = model("image.jpg")
|
1021
|
+
>>> boxes = results[0].boxes
|
1022
|
+
>>> xyxy = boxes.xyxy
|
1023
|
+
>>> print(xyxy)
|
1024
|
+
"""
|
468
1025
|
return self.data[:, :4]
|
469
1026
|
|
470
1027
|
@property
|
471
1028
|
def conf(self):
|
472
|
-
"""
|
1029
|
+
"""
|
1030
|
+
Returns the confidence scores for each detection box.
|
1031
|
+
|
1032
|
+
Returns:
|
1033
|
+
(torch.Tensor | numpy.ndarray): A 1D tensor or array containing confidence scores for each detection,
|
1034
|
+
with shape (N,) where N is the number of detections.
|
1035
|
+
|
1036
|
+
Examples:
|
1037
|
+
>>> boxes = Boxes(torch.tensor([[10, 20, 30, 40, 0.9, 0]]), orig_shape=(100, 100))
|
1038
|
+
>>> conf_scores = boxes.conf
|
1039
|
+
>>> print(conf_scores)
|
1040
|
+
tensor([0.9000])
|
1041
|
+
"""
|
473
1042
|
return self.data[:, -2]
|
474
1043
|
|
475
1044
|
@property
|
476
1045
|
def cls(self):
|
477
|
-
"""
|
1046
|
+
"""
|
1047
|
+
Returns the class ID tensor representing category predictions for each bounding box.
|
1048
|
+
|
1049
|
+
Returns:
|
1050
|
+
(torch.Tensor | numpy.ndarray): A tensor or numpy array containing the class IDs for each detection box.
|
1051
|
+
The shape is (N,), where N is the number of boxes.
|
1052
|
+
|
1053
|
+
Examples:
|
1054
|
+
>>> results = model("image.jpg")
|
1055
|
+
>>> boxes = results[0].boxes
|
1056
|
+
>>> class_ids = boxes.cls
|
1057
|
+
>>> print(class_ids) # tensor([0., 2., 1.])
|
1058
|
+
"""
|
478
1059
|
return self.data[:, -1]
|
479
1060
|
|
480
1061
|
@property
|
481
1062
|
def id(self):
|
482
|
-
"""
|
1063
|
+
"""
|
1064
|
+
Returns the tracking IDs for each detection box if available.
|
1065
|
+
|
1066
|
+
Returns:
|
1067
|
+
(torch.Tensor | None): A tensor containing tracking IDs for each box if tracking is enabled,
|
1068
|
+
otherwise None. Shape is (N,) where N is the number of boxes.
|
1069
|
+
|
1070
|
+
Examples:
|
1071
|
+
>>> results = model.track("path/to/video.mp4")
|
1072
|
+
>>> for result in results:
|
1073
|
+
... boxes = result.boxes
|
1074
|
+
... if boxes.is_track:
|
1075
|
+
... track_ids = boxes.id
|
1076
|
+
... print(f"Tracking IDs: {track_ids}")
|
1077
|
+
... else:
|
1078
|
+
... print("Tracking is not enabled for these boxes.")
|
1079
|
+
|
1080
|
+
Notes:
|
1081
|
+
- This property is only available when tracking is enabled (i.e., when `is_track` is True).
|
1082
|
+
- The tracking IDs are typically used to associate detections across multiple frames in video analysis.
|
1083
|
+
"""
|
483
1084
|
return self.data[:, -3] if self.is_track else None
|
484
1085
|
|
485
1086
|
@property
|
486
1087
|
@lru_cache(maxsize=2) # maxsize 1 should suffice
|
487
1088
|
def xywh(self):
|
488
|
-
"""
|
1089
|
+
"""
|
1090
|
+
Convert bounding boxes from [x1, y1, x2, y2] format to [x, y, width, height] format.
|
1091
|
+
|
1092
|
+
Returns:
|
1093
|
+
(torch.Tensor | numpy.ndarray): Boxes in [x_center, y_center, width, height] format, where x_center, y_center are the coordinates of
|
1094
|
+
the center point of the bounding box, width, height are the dimensions of the bounding box and the
|
1095
|
+
shape of the returned tensor is (N, 4), where N is the number of boxes.
|
1096
|
+
|
1097
|
+
Examples:
|
1098
|
+
>>> boxes = Boxes(torch.tensor([[100, 50, 150, 100], [200, 150, 300, 250]]), orig_shape=(480, 640))
|
1099
|
+
>>> xywh = boxes.xywh
|
1100
|
+
>>> print(xywh)
|
1101
|
+
tensor([[100.0000, 50.0000, 50.0000, 50.0000],
|
1102
|
+
[200.0000, 150.0000, 100.0000, 100.0000]])
|
1103
|
+
"""
|
489
1104
|
return ops.xyxy2xywh(self.xyxy)
|
490
1105
|
|
491
1106
|
@property
|
492
1107
|
@lru_cache(maxsize=2)
|
493
1108
|
def xyxyn(self):
|
494
|
-
"""
|
1109
|
+
"""
|
1110
|
+
Returns normalized bounding box coordinates relative to the original image size.
|
1111
|
+
|
1112
|
+
This property calculates and returns the bounding box coordinates in [x1, y1, x2, y2] format,
|
1113
|
+
normalized to the range [0, 1] based on the original image dimensions.
|
1114
|
+
|
1115
|
+
Returns:
|
1116
|
+
(torch.Tensor | numpy.ndarray): Normalized bounding box coordinates with shape (N, 4), where N is
|
1117
|
+
the number of boxes. Each row contains [x1, y1, x2, y2] values normalized to [0, 1].
|
1118
|
+
|
1119
|
+
Examples:
|
1120
|
+
>>> boxes = Boxes(torch.tensor([[100, 50, 300, 400, 0.9, 0]]), orig_shape=(480, 640))
|
1121
|
+
>>> normalized = boxes.xyxyn
|
1122
|
+
>>> print(normalized)
|
1123
|
+
tensor([[0.1562, 0.1042, 0.4688, 0.8333]])
|
1124
|
+
"""
|
495
1125
|
xyxy = self.xyxy.clone() if isinstance(self.xyxy, torch.Tensor) else np.copy(self.xyxy)
|
496
1126
|
xyxy[..., [0, 2]] /= self.orig_shape[1]
|
497
1127
|
xyxy[..., [1, 3]] /= self.orig_shape[0]
|
@@ -500,7 +1130,23 @@ class Boxes(BaseTensor):
|
|
500
1130
|
@property
|
501
1131
|
@lru_cache(maxsize=2)
|
502
1132
|
def xywhn(self):
|
503
|
-
"""
|
1133
|
+
"""
|
1134
|
+
Returns normalized bounding boxes in [x, y, width, height] format.
|
1135
|
+
|
1136
|
+
This property calculates and returns the normalized bounding box coordinates in the format
|
1137
|
+
[x_center, y_center, width, height], where all values are relative to the original image dimensions.
|
1138
|
+
|
1139
|
+
Returns:
|
1140
|
+
(torch.Tensor | numpy.ndarray): Normalized bounding boxes with shape (N, 4), where N is the
|
1141
|
+
number of boxes. Each row contains [x_center, y_center, width, height] values normalized
|
1142
|
+
to [0, 1] based on the original image dimensions.
|
1143
|
+
|
1144
|
+
Examples:
|
1145
|
+
>>> boxes = Boxes(torch.tensor([[100, 50, 150, 100, 0.9, 0]]), orig_shape=(480, 640))
|
1146
|
+
>>> normalized = boxes.xywhn
|
1147
|
+
>>> print(normalized)
|
1148
|
+
tensor([[0.1953, 0.1562, 0.0781, 0.1042]])
|
1149
|
+
"""
|
504
1150
|
xywh = ops.xyxy2xywh(self.xyxy)
|
505
1151
|
xywh[..., [0, 2]] /= self.orig_shape[1]
|
506
1152
|
xywh[..., [1, 3]] /= self.orig_shape[0]
|
@@ -511,19 +1157,44 @@ class Masks(BaseTensor):
|
|
511
1157
|
"""
|
512
1158
|
A class for storing and manipulating detection masks.
|
513
1159
|
|
1160
|
+
This class extends BaseTensor and provides functionality for handling segmentation masks,
|
1161
|
+
including methods for converting between pixel and normalized coordinates.
|
1162
|
+
|
514
1163
|
Attributes:
|
515
|
-
|
516
|
-
|
1164
|
+
data (torch.Tensor | numpy.ndarray): The raw tensor or array containing mask data.
|
1165
|
+
orig_shape (tuple): Original image shape in (height, width) format.
|
1166
|
+
xy (List[numpy.ndarray]): A list of segments in pixel coordinates.
|
1167
|
+
xyn (List[numpy.ndarray]): A list of normalized segments.
|
517
1168
|
|
518
1169
|
Methods:
|
519
|
-
cpu(): Returns the
|
520
|
-
numpy(): Returns the
|
521
|
-
cuda(): Returns the
|
522
|
-
to(
|
1170
|
+
cpu(): Returns a copy of the Masks object with the mask tensor on CPU memory.
|
1171
|
+
numpy(): Returns a copy of the Masks object with the mask tensor as a numpy array.
|
1172
|
+
cuda(): Returns a copy of the Masks object with the mask tensor on GPU memory.
|
1173
|
+
to(*args, **kwargs): Returns a copy of the Masks object with the mask tensor on specified device and dtype.
|
1174
|
+
|
1175
|
+
Examples:
|
1176
|
+
>>> masks_data = torch.rand(1, 160, 160)
|
1177
|
+
>>> orig_shape = (720, 1280)
|
1178
|
+
>>> masks = Masks(masks_data, orig_shape)
|
1179
|
+
>>> pixel_coords = masks.xy
|
1180
|
+
>>> normalized_coords = masks.xyn
|
523
1181
|
"""
|
524
1182
|
|
525
1183
|
def __init__(self, masks, orig_shape) -> None:
|
526
|
-
"""
|
1184
|
+
"""
|
1185
|
+
Initialize the Masks class with detection mask data and the original image shape.
|
1186
|
+
|
1187
|
+
Args:
|
1188
|
+
masks (torch.Tensor | np.ndarray): Detection masks with shape (num_masks, height, width).
|
1189
|
+
orig_shape (tuple): The original image shape as (height, width). Used for normalization.
|
1190
|
+
|
1191
|
+
Examples:
|
1192
|
+
>>> import torch
|
1193
|
+
>>> from ultralytics.engine.results import Masks
|
1194
|
+
>>> masks = torch.rand(10, 160, 160) # 10 masks of 160x160 resolution
|
1195
|
+
>>> orig_shape = (720, 1280) # Original image shape
|
1196
|
+
>>> mask_obj = Masks(masks, orig_shape)
|
1197
|
+
"""
|
527
1198
|
if masks.ndim == 2:
|
528
1199
|
masks = masks[None, :]
|
529
1200
|
super().__init__(masks, orig_shape)
|
@@ -531,7 +1202,23 @@ class Masks(BaseTensor):
|
|
531
1202
|
@property
|
532
1203
|
@lru_cache(maxsize=1)
|
533
1204
|
def xyn(self):
|
534
|
-
"""
|
1205
|
+
"""
|
1206
|
+
Returns normalized xy-coordinates of the segmentation masks.
|
1207
|
+
|
1208
|
+
This property calculates and caches the normalized xy-coordinates of the segmentation masks. The coordinates
|
1209
|
+
are normalized relative to the original image shape.
|
1210
|
+
|
1211
|
+
Returns:
|
1212
|
+
(List[numpy.ndarray]): A list of numpy arrays, where each array contains the normalized xy-coordinates
|
1213
|
+
of a single segmentation mask. Each array has shape (N, 2), where N is the number of points in the
|
1214
|
+
mask contour.
|
1215
|
+
|
1216
|
+
Examples:
|
1217
|
+
>>> results = model("image.jpg")
|
1218
|
+
>>> masks = results[0].masks
|
1219
|
+
>>> normalized_coords = masks.xyn
|
1220
|
+
>>> print(normalized_coords[0]) # Normalized coordinates of the first mask
|
1221
|
+
"""
|
535
1222
|
return [
|
536
1223
|
ops.scale_coords(self.data.shape[1:], x, self.orig_shape, normalize=True)
|
537
1224
|
for x in ops.masks2segments(self.data)
|
@@ -540,7 +1227,24 @@ class Masks(BaseTensor):
|
|
540
1227
|
@property
|
541
1228
|
@lru_cache(maxsize=1)
|
542
1229
|
def xy(self):
|
543
|
-
"""
|
1230
|
+
"""
|
1231
|
+
Returns the [x, y] pixel coordinates for each segment in the mask tensor.
|
1232
|
+
|
1233
|
+
This property calculates and returns a list of pixel coordinates for each segmentation mask in the
|
1234
|
+
Masks object. The coordinates are scaled to match the original image dimensions.
|
1235
|
+
|
1236
|
+
Returns:
|
1237
|
+
(List[numpy.ndarray]): A list of numpy arrays, where each array contains the [x, y] pixel
|
1238
|
+
coordinates for a single segmentation mask. Each array has shape (N, 2), where N is the
|
1239
|
+
number of points in the segment.
|
1240
|
+
|
1241
|
+
Examples:
|
1242
|
+
>>> results = model("image.jpg")
|
1243
|
+
>>> masks = results[0].masks
|
1244
|
+
>>> xy_coords = masks.xy
|
1245
|
+
>>> print(len(xy_coords)) # Number of masks
|
1246
|
+
>>> print(xy_coords[0].shape) # Shape of first mask's coordinates
|
1247
|
+
"""
|
544
1248
|
return [
|
545
1249
|
ops.scale_coords(self.data.shape[1:], x, self.orig_shape, normalize=False)
|
546
1250
|
for x in ops.masks2segments(self.data)
|
@@ -551,21 +1255,53 @@ class Keypoints(BaseTensor):
|
|
551
1255
|
"""
|
552
1256
|
A class for storing and manipulating detection keypoints.
|
553
1257
|
|
1258
|
+
This class encapsulates functionality for handling keypoint data, including coordinate manipulation,
|
1259
|
+
normalization, and confidence values.
|
1260
|
+
|
554
1261
|
Attributes:
|
555
|
-
|
556
|
-
|
557
|
-
|
1262
|
+
data (torch.Tensor): The raw tensor containing keypoint data.
|
1263
|
+
orig_shape (Tuple[int, int]): The original image dimensions (height, width).
|
1264
|
+
has_visible (bool): Indicates whether visibility information is available for keypoints.
|
1265
|
+
xy (torch.Tensor): Keypoint coordinates in [x, y] format.
|
1266
|
+
xyn (torch.Tensor): Normalized keypoint coordinates in [x, y] format, relative to orig_shape.
|
1267
|
+
conf (torch.Tensor): Confidence values for each keypoint, if available.
|
558
1268
|
|
559
1269
|
Methods:
|
560
1270
|
cpu(): Returns a copy of the keypoints tensor on CPU memory.
|
561
1271
|
numpy(): Returns a copy of the keypoints tensor as a numpy array.
|
562
1272
|
cuda(): Returns a copy of the keypoints tensor on GPU memory.
|
563
|
-
to(
|
1273
|
+
to(*args, **kwargs): Returns a copy of the keypoints tensor with specified device and dtype.
|
1274
|
+
|
1275
|
+
Examples:
|
1276
|
+
>>> import torch
|
1277
|
+
>>> from ultralytics.engine.results import Keypoints
|
1278
|
+
>>> keypoints_data = torch.rand(1, 17, 3) # 1 detection, 17 keypoints, (x, y, conf)
|
1279
|
+
>>> orig_shape = (480, 640) # Original image shape (height, width)
|
1280
|
+
>>> keypoints = Keypoints(keypoints_data, orig_shape)
|
1281
|
+
>>> print(keypoints.xy.shape) # Access xy coordinates
|
1282
|
+
>>> print(keypoints.conf) # Access confidence values
|
1283
|
+
>>> keypoints_cpu = keypoints.cpu() # Move keypoints to CPU
|
564
1284
|
"""
|
565
1285
|
|
566
1286
|
@smart_inference_mode() # avoid keypoints < conf in-place error
|
567
1287
|
def __init__(self, keypoints, orig_shape) -> None:
|
568
|
-
"""
|
1288
|
+
"""
|
1289
|
+
Initializes the Keypoints object with detection keypoints and original image dimensions.
|
1290
|
+
|
1291
|
+
This method processes the input keypoints tensor, handling both 2D and 3D formats. For 3D tensors
|
1292
|
+
(x, y, confidence), it masks out low-confidence keypoints by setting their coordinates to zero.
|
1293
|
+
|
1294
|
+
Args:
|
1295
|
+
keypoints (torch.Tensor): A tensor containing keypoint data. Shape can be either:
|
1296
|
+
- (num_objects, num_keypoints, 2) for x, y coordinates only
|
1297
|
+
- (num_objects, num_keypoints, 3) for x, y coordinates and confidence scores
|
1298
|
+
orig_shape (Tuple[int, int]): The original image dimensions (height, width).
|
1299
|
+
|
1300
|
+
Examples:
|
1301
|
+
>>> kpts = torch.rand(1, 17, 3) # 1 object, 17 keypoints (COCO format), x,y,conf
|
1302
|
+
>>> orig_shape = (720, 1280) # Original image height, width
|
1303
|
+
>>> keypoints = Keypoints(kpts, orig_shape)
|
1304
|
+
"""
|
569
1305
|
if keypoints.ndim == 2:
|
570
1306
|
keypoints = keypoints[None, :]
|
571
1307
|
if keypoints.shape[2] == 3: # x, y, conf
|
@@ -577,13 +1313,44 @@ class Keypoints(BaseTensor):
|
|
577
1313
|
@property
|
578
1314
|
@lru_cache(maxsize=1)
|
579
1315
|
def xy(self):
|
580
|
-
"""
|
1316
|
+
"""
|
1317
|
+
Returns x, y coordinates of keypoints.
|
1318
|
+
|
1319
|
+
Returns:
|
1320
|
+
(torch.Tensor): A tensor containing the x, y coordinates of keypoints with shape (N, K, 2), where N is
|
1321
|
+
the number of detections and K is the number of keypoints per detection.
|
1322
|
+
|
1323
|
+
Examples:
|
1324
|
+
>>> results = model("image.jpg")
|
1325
|
+
>>> keypoints = results[0].keypoints
|
1326
|
+
>>> xy = keypoints.xy
|
1327
|
+
>>> print(xy.shape) # (N, K, 2)
|
1328
|
+
>>> print(xy[0]) # x, y coordinates of keypoints for first detection
|
1329
|
+
|
1330
|
+
Notes:
|
1331
|
+
- The returned coordinates are in pixel units relative to the original image dimensions.
|
1332
|
+
- If keypoints were initialized with confidence values, only keypoints with confidence >= 0.5 are returned.
|
1333
|
+
- This property uses LRU caching to improve performance on repeated access.
|
1334
|
+
"""
|
581
1335
|
return self.data[..., :2]
|
582
1336
|
|
583
1337
|
@property
|
584
1338
|
@lru_cache(maxsize=1)
|
585
1339
|
def xyn(self):
|
586
|
-
"""
|
1340
|
+
"""
|
1341
|
+
Returns normalized coordinates (x, y) of keypoints relative to the original image size.
|
1342
|
+
|
1343
|
+
Returns:
|
1344
|
+
(torch.Tensor | numpy.ndarray): A tensor or array of shape (N, K, 2) containing normalized keypoint
|
1345
|
+
coordinates, where N is the number of instances, K is the number of keypoints, and the last
|
1346
|
+
dimension contains [x, y] values in the range [0, 1].
|
1347
|
+
|
1348
|
+
Examples:
|
1349
|
+
>>> keypoints = Keypoints(torch.rand(1, 17, 2), orig_shape=(480, 640))
|
1350
|
+
>>> normalized_kpts = keypoints.xyn
|
1351
|
+
>>> print(normalized_kpts.shape)
|
1352
|
+
torch.Size([1, 17, 2])
|
1353
|
+
"""
|
587
1354
|
xy = self.xy.clone() if isinstance(self.xy, torch.Tensor) else np.copy(self.xy)
|
588
1355
|
xy[..., 0] /= self.orig_shape[1]
|
589
1356
|
xy[..., 1] /= self.orig_shape[0]
|
@@ -592,53 +1359,160 @@ class Keypoints(BaseTensor):
|
|
592
1359
|
@property
|
593
1360
|
@lru_cache(maxsize=1)
|
594
1361
|
def conf(self):
|
595
|
-
"""
|
1362
|
+
"""
|
1363
|
+
Returns confidence values for each keypoint.
|
1364
|
+
|
1365
|
+
Returns:
|
1366
|
+
(torch.Tensor | None): A tensor containing confidence scores for each keypoint if available,
|
1367
|
+
otherwise None. Shape is (num_detections, num_keypoints) for batched data or (num_keypoints,)
|
1368
|
+
for single detection.
|
1369
|
+
|
1370
|
+
Examples:
|
1371
|
+
>>> keypoints = Keypoints(torch.rand(1, 17, 3), orig_shape=(640, 640)) # 1 detection, 17 keypoints
|
1372
|
+
>>> conf = keypoints.conf
|
1373
|
+
>>> print(conf.shape) # torch.Size([1, 17])
|
1374
|
+
"""
|
596
1375
|
return self.data[..., 2] if self.has_visible else None
|
597
1376
|
|
598
1377
|
|
599
1378
|
class Probs(BaseTensor):
|
600
1379
|
"""
|
601
|
-
A class for storing and manipulating classification
|
1380
|
+
A class for storing and manipulating classification probabilities.
|
1381
|
+
|
1382
|
+
This class extends BaseTensor and provides methods for accessing and manipulating
|
1383
|
+
classification probabilities, including top-1 and top-5 predictions.
|
602
1384
|
|
603
1385
|
Attributes:
|
604
|
-
|
605
|
-
|
606
|
-
|
607
|
-
|
1386
|
+
data (torch.Tensor | numpy.ndarray): The raw tensor or array containing classification probabilities.
|
1387
|
+
orig_shape (tuple | None): The original image shape as (height, width). Not used in this class.
|
1388
|
+
top1 (int): Index of the class with the highest probability.
|
1389
|
+
top5 (List[int]): Indices of the top 5 classes by probability.
|
1390
|
+
top1conf (torch.Tensor | numpy.ndarray): Confidence score of the top 1 class.
|
1391
|
+
top5conf (torch.Tensor | numpy.ndarray): Confidence scores of the top 5 classes.
|
608
1392
|
|
609
1393
|
Methods:
|
610
|
-
cpu(): Returns a copy of the
|
611
|
-
numpy(): Returns a copy of the
|
612
|
-
cuda(): Returns a copy of the
|
613
|
-
to(): Returns a copy of the
|
1394
|
+
cpu(): Returns a copy of the probabilities tensor on CPU memory.
|
1395
|
+
numpy(): Returns a copy of the probabilities tensor as a numpy array.
|
1396
|
+
cuda(): Returns a copy of the probabilities tensor on GPU memory.
|
1397
|
+
to(*args, **kwargs): Returns a copy of the probabilities tensor with specified device and dtype.
|
1398
|
+
|
1399
|
+
Examples:
|
1400
|
+
>>> probs = torch.tensor([0.1, 0.3, 0.6])
|
1401
|
+
>>> p = Probs(probs)
|
1402
|
+
>>> print(p.top1)
|
1403
|
+
2
|
1404
|
+
>>> print(p.top5)
|
1405
|
+
[2, 1, 0]
|
1406
|
+
>>> print(p.top1conf)
|
1407
|
+
tensor(0.6000)
|
1408
|
+
>>> print(p.top5conf)
|
1409
|
+
tensor([0.6000, 0.3000, 0.1000])
|
614
1410
|
"""
|
615
1411
|
|
616
1412
|
def __init__(self, probs, orig_shape=None) -> None:
|
617
|
-
"""
|
1413
|
+
"""
|
1414
|
+
Initialize the Probs class with classification probabilities.
|
1415
|
+
|
1416
|
+
This class stores and manages classification probabilities, providing easy access to top predictions and their
|
1417
|
+
confidences.
|
1418
|
+
|
1419
|
+
Args:
|
1420
|
+
probs (torch.Tensor | np.ndarray): A 1D tensor or array of classification probabilities.
|
1421
|
+
orig_shape (tuple | None): The original image shape as (height, width). Not used in this class but kept for
|
1422
|
+
consistency with other result classes.
|
1423
|
+
|
1424
|
+
Attributes:
|
1425
|
+
data (torch.Tensor | np.ndarray): The raw tensor or array containing classification probabilities.
|
1426
|
+
top1 (int): Index of the top 1 class.
|
1427
|
+
top5 (List[int]): Indices of the top 5 classes.
|
1428
|
+
top1conf (torch.Tensor | np.ndarray): Confidence of the top 1 class.
|
1429
|
+
top5conf (torch.Tensor | np.ndarray): Confidences of the top 5 classes.
|
1430
|
+
|
1431
|
+
Examples:
|
1432
|
+
>>> import torch
|
1433
|
+
>>> probs = torch.tensor([0.1, 0.3, 0.2, 0.4])
|
1434
|
+
>>> p = Probs(probs)
|
1435
|
+
>>> print(p.top1)
|
1436
|
+
3
|
1437
|
+
>>> print(p.top1conf)
|
1438
|
+
tensor(0.4000)
|
1439
|
+
>>> print(p.top5)
|
1440
|
+
[3, 1, 2, 0]
|
1441
|
+
"""
|
618
1442
|
super().__init__(probs, orig_shape)
|
619
1443
|
|
620
1444
|
@property
|
621
1445
|
@lru_cache(maxsize=1)
|
622
1446
|
def top1(self):
|
623
|
-
"""
|
1447
|
+
"""
|
1448
|
+
Returns the index of the class with the highest probability.
|
1449
|
+
|
1450
|
+
Returns:
|
1451
|
+
(int): Index of the class with the highest probability.
|
1452
|
+
|
1453
|
+
Examples:
|
1454
|
+
>>> probs = Probs(torch.tensor([0.1, 0.3, 0.6]))
|
1455
|
+
>>> probs.top1
|
1456
|
+
2
|
1457
|
+
"""
|
624
1458
|
return int(self.data.argmax())
|
625
1459
|
|
626
1460
|
@property
|
627
1461
|
@lru_cache(maxsize=1)
|
628
1462
|
def top5(self):
|
629
|
-
"""
|
1463
|
+
"""
|
1464
|
+
Returns the indices of the top 5 class probabilities.
|
1465
|
+
|
1466
|
+
Returns:
|
1467
|
+
(List[int]): A list containing the indices of the top 5 class probabilities, sorted in descending order.
|
1468
|
+
|
1469
|
+
Examples:
|
1470
|
+
>>> probs = Probs(torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5]))
|
1471
|
+
>>> print(probs.top5)
|
1472
|
+
[4, 3, 2, 1, 0]
|
1473
|
+
"""
|
630
1474
|
return (-self.data).argsort(0)[:5].tolist() # this way works with both torch and numpy.
|
631
1475
|
|
632
1476
|
@property
|
633
1477
|
@lru_cache(maxsize=1)
|
634
1478
|
def top1conf(self):
|
635
|
-
"""
|
1479
|
+
"""
|
1480
|
+
Returns the confidence score of the highest probability class.
|
1481
|
+
|
1482
|
+
This property retrieves the confidence score (probability) of the class with the highest predicted probability
|
1483
|
+
from the classification results.
|
1484
|
+
|
1485
|
+
Returns:
|
1486
|
+
(torch.Tensor | numpy.ndarray): A tensor containing the confidence score of the top 1 class.
|
1487
|
+
|
1488
|
+
Examples:
|
1489
|
+
>>> results = model("image.jpg") # classify an image
|
1490
|
+
>>> probs = results[0].probs # get classification probabilities
|
1491
|
+
>>> top1_confidence = probs.top1conf # get confidence of top 1 class
|
1492
|
+
>>> print(f"Top 1 class confidence: {top1_confidence.item():.4f}")
|
1493
|
+
"""
|
636
1494
|
return self.data[self.top1]
|
637
1495
|
|
638
1496
|
@property
|
639
1497
|
@lru_cache(maxsize=1)
|
640
1498
|
def top5conf(self):
|
641
|
-
"""
|
1499
|
+
"""
|
1500
|
+
Returns confidence scores for the top 5 classification predictions.
|
1501
|
+
|
1502
|
+
This property retrieves the confidence scores corresponding to the top 5 class probabilities
|
1503
|
+
predicted by the model. It provides a quick way to access the most likely class predictions
|
1504
|
+
along with their associated confidence levels.
|
1505
|
+
|
1506
|
+
Returns:
|
1507
|
+
(torch.Tensor | numpy.ndarray): A tensor or array containing the confidence scores for the
|
1508
|
+
top 5 predicted classes, sorted in descending order of probability.
|
1509
|
+
|
1510
|
+
Examples:
|
1511
|
+
>>> results = model("image.jpg")
|
1512
|
+
>>> probs = results[0].probs
|
1513
|
+
>>> top5_conf = probs.top5conf
|
1514
|
+
>>> print(top5_conf) # Prints confidence scores for top 5 classes
|
1515
|
+
"""
|
642
1516
|
return self.data[self.top5]
|
643
1517
|
|
644
1518
|
|
@@ -646,69 +1520,182 @@ class OBB(BaseTensor):
|
|
646
1520
|
"""
|
647
1521
|
A class for storing and manipulating Oriented Bounding Boxes (OBB).
|
648
1522
|
|
649
|
-
|
650
|
-
|
651
|
-
with shape (num_boxes, 7) or (num_boxes, 8). The last two columns contain confidence and class values.
|
652
|
-
If present, the third last column contains track IDs, and the fifth column from the left contains rotation.
|
653
|
-
orig_shape (tuple): Original image size, in the format (height, width).
|
1523
|
+
This class provides functionality to handle oriented bounding boxes, including conversion between
|
1524
|
+
different formats, normalization, and access to various properties of the boxes.
|
654
1525
|
|
655
1526
|
Attributes:
|
656
|
-
|
657
|
-
|
658
|
-
|
659
|
-
|
660
|
-
|
661
|
-
|
662
|
-
|
663
|
-
|
1527
|
+
data (torch.Tensor): The raw OBB tensor containing box coordinates and associated data.
|
1528
|
+
orig_shape (tuple): Original image size as (height, width).
|
1529
|
+
is_track (bool): Indicates whether tracking IDs are included in the box data.
|
1530
|
+
xywhr (torch.Tensor | numpy.ndarray): Boxes in [x_center, y_center, width, height, rotation] format.
|
1531
|
+
conf (torch.Tensor | numpy.ndarray): Confidence scores for each box.
|
1532
|
+
cls (torch.Tensor | numpy.ndarray): Class labels for each box.
|
1533
|
+
id (torch.Tensor | numpy.ndarray): Tracking IDs for each box, if available.
|
1534
|
+
xyxyxyxy (torch.Tensor | numpy.ndarray): Boxes in 8-point [x1, y1, x2, y2, x3, y3, x4, y4] format.
|
1535
|
+
xyxyxyxyn (torch.Tensor | numpy.ndarray): Normalized 8-point coordinates relative to orig_shape.
|
1536
|
+
xyxy (torch.Tensor | numpy.ndarray): Axis-aligned bounding boxes in [x1, y1, x2, y2] format.
|
664
1537
|
|
665
1538
|
Methods:
|
666
|
-
cpu():
|
667
|
-
numpy():
|
668
|
-
cuda():
|
669
|
-
to(*args, **kwargs):
|
1539
|
+
cpu(): Returns a copy of the OBB object with all tensors on CPU memory.
|
1540
|
+
numpy(): Returns a copy of the OBB object with all tensors as numpy arrays.
|
1541
|
+
cuda(): Returns a copy of the OBB object with all tensors on GPU memory.
|
1542
|
+
to(*args, **kwargs): Returns a copy of the OBB object with tensors on specified device and dtype.
|
1543
|
+
|
1544
|
+
Examples:
|
1545
|
+
>>> boxes = torch.tensor([[100, 50, 150, 100, 30, 0.9, 0]]) # xywhr, conf, cls
|
1546
|
+
>>> obb = OBB(boxes, orig_shape=(480, 640))
|
1547
|
+
>>> print(obb.xyxyxyxy)
|
1548
|
+
>>> print(obb.conf)
|
1549
|
+
>>> print(obb.cls)
|
670
1550
|
"""
|
671
1551
|
|
672
1552
|
def __init__(self, boxes, orig_shape) -> None:
|
673
|
-
"""
|
1553
|
+
"""
|
1554
|
+
Initialize an OBB (Oriented Bounding Box) instance with oriented bounding box data and original image shape.
|
1555
|
+
|
1556
|
+
This class stores and manipulates Oriented Bounding Boxes (OBB) for object detection tasks. It provides
|
1557
|
+
various properties and methods to access and transform the OBB data.
|
1558
|
+
|
1559
|
+
Args:
|
1560
|
+
boxes (torch.Tensor | numpy.ndarray): A tensor or numpy array containing the detection boxes,
|
1561
|
+
with shape (num_boxes, 7) or (num_boxes, 8). The last two columns contain confidence and class values.
|
1562
|
+
If present, the third last column contains track IDs, and the fifth column contains rotation.
|
1563
|
+
orig_shape (Tuple[int, int]): Original image size, in the format (height, width).
|
1564
|
+
|
1565
|
+
Attributes:
|
1566
|
+
data (torch.Tensor | numpy.ndarray): The raw OBB tensor.
|
1567
|
+
orig_shape (Tuple[int, int]): The original image shape.
|
1568
|
+
is_track (bool): Whether the boxes include tracking IDs.
|
1569
|
+
|
1570
|
+
Raises:
|
1571
|
+
AssertionError: If the number of values per box is not 7 or 8.
|
1572
|
+
|
1573
|
+
Examples:
|
1574
|
+
>>> import torch
|
1575
|
+
>>> boxes = torch.rand(3, 7) # 3 boxes with 7 values each
|
1576
|
+
>>> orig_shape = (640, 480)
|
1577
|
+
>>> obb = OBB(boxes, orig_shape)
|
1578
|
+
>>> print(obb.xywhr) # Access the boxes in xywhr format
|
1579
|
+
"""
|
674
1580
|
if boxes.ndim == 1:
|
675
1581
|
boxes = boxes[None, :]
|
676
1582
|
n = boxes.shape[-1]
|
677
|
-
assert n in
|
1583
|
+
assert n in {7, 8}, f"expected 7 or 8 values but got {n}" # xywh, rotation, track_id, conf, cls
|
678
1584
|
super().__init__(boxes, orig_shape)
|
679
1585
|
self.is_track = n == 8
|
680
1586
|
self.orig_shape = orig_shape
|
681
1587
|
|
682
1588
|
@property
|
683
1589
|
def xywhr(self):
|
684
|
-
"""
|
1590
|
+
"""
|
1591
|
+
Returns boxes in [x_center, y_center, width, height, rotation] format.
|
1592
|
+
|
1593
|
+
Returns:
|
1594
|
+
(torch.Tensor | numpy.ndarray): A tensor or numpy array containing the oriented bounding boxes with format
|
1595
|
+
[x_center, y_center, width, height, rotation]. The shape is (N, 5) where N is the number of boxes.
|
1596
|
+
|
1597
|
+
Examples:
|
1598
|
+
>>> results = model("image.jpg")
|
1599
|
+
>>> obb = results[0].obb
|
1600
|
+
>>> xywhr = obb.xywhr
|
1601
|
+
>>> print(xywhr.shape)
|
1602
|
+
torch.Size([3, 5])
|
1603
|
+
"""
|
685
1604
|
return self.data[:, :5]
|
686
1605
|
|
687
1606
|
@property
|
688
1607
|
def conf(self):
|
689
|
-
"""
|
1608
|
+
"""
|
1609
|
+
Returns the confidence scores for Oriented Bounding Boxes (OBBs).
|
1610
|
+
|
1611
|
+
This property retrieves the confidence values associated with each OBB detection. The confidence score
|
1612
|
+
represents the model's certainty in the detection.
|
1613
|
+
|
1614
|
+
Returns:
|
1615
|
+
(torch.Tensor | numpy.ndarray): A tensor or numpy array of shape (N,) containing confidence scores
|
1616
|
+
for N detections, where each score is in the range [0, 1].
|
1617
|
+
|
1618
|
+
Examples:
|
1619
|
+
>>> results = model("image.jpg")
|
1620
|
+
>>> obb_result = results[0].obb
|
1621
|
+
>>> confidence_scores = obb_result.conf
|
1622
|
+
>>> print(confidence_scores)
|
1623
|
+
"""
|
690
1624
|
return self.data[:, -2]
|
691
1625
|
|
692
1626
|
@property
|
693
1627
|
def cls(self):
|
694
|
-
"""
|
1628
|
+
"""
|
1629
|
+
Returns the class values of the oriented bounding boxes.
|
1630
|
+
|
1631
|
+
Returns:
|
1632
|
+
(torch.Tensor | numpy.ndarray): A tensor or numpy array containing the class values for each oriented
|
1633
|
+
bounding box. The shape is (N,), where N is the number of boxes.
|
1634
|
+
|
1635
|
+
Examples:
|
1636
|
+
>>> results = model("image.jpg")
|
1637
|
+
>>> result = results[0]
|
1638
|
+
>>> obb = result.obb
|
1639
|
+
>>> class_values = obb.cls
|
1640
|
+
>>> print(class_values)
|
1641
|
+
"""
|
695
1642
|
return self.data[:, -1]
|
696
1643
|
|
697
1644
|
@property
|
698
1645
|
def id(self):
|
699
|
-
"""
|
1646
|
+
"""
|
1647
|
+
Returns the tracking IDs of the oriented bounding boxes (if available).
|
1648
|
+
|
1649
|
+
Returns:
|
1650
|
+
(torch.Tensor | numpy.ndarray | None): A tensor or numpy array containing the tracking IDs for each
|
1651
|
+
oriented bounding box. Returns None if tracking IDs are not available.
|
1652
|
+
|
1653
|
+
Examples:
|
1654
|
+
>>> results = model("image.jpg", tracker=True) # Run inference with tracking
|
1655
|
+
>>> for result in results:
|
1656
|
+
... if result.obb is not None:
|
1657
|
+
... track_ids = result.obb.id
|
1658
|
+
... if track_ids is not None:
|
1659
|
+
... print(f"Tracking IDs: {track_ids}")
|
1660
|
+
"""
|
700
1661
|
return self.data[:, -3] if self.is_track else None
|
701
1662
|
|
702
1663
|
@property
|
703
1664
|
@lru_cache(maxsize=2)
|
704
1665
|
def xyxyxyxy(self):
|
705
|
-
"""
|
1666
|
+
"""
|
1667
|
+
Converts OBB format to 8-point (xyxyxyxy) coordinate format for rotated bounding boxes.
|
1668
|
+
|
1669
|
+
Returns:
|
1670
|
+
(torch.Tensor | numpy.ndarray): Rotated bounding boxes in xyxyxyxy format with shape (N, 4, 2), where N is
|
1671
|
+
the number of boxes. Each box is represented by 4 points (x, y), starting from the top-left corner and
|
1672
|
+
moving clockwise.
|
1673
|
+
|
1674
|
+
Examples:
|
1675
|
+
>>> obb = OBB(torch.tensor([[100, 100, 50, 30, 0.5, 0.9, 0]]), orig_shape=(640, 640))
|
1676
|
+
>>> xyxyxyxy = obb.xyxyxyxy
|
1677
|
+
>>> print(xyxyxyxy.shape)
|
1678
|
+
torch.Size([1, 4, 2])
|
1679
|
+
"""
|
706
1680
|
return ops.xywhr2xyxyxyxy(self.xywhr)
|
707
1681
|
|
708
1682
|
@property
|
709
1683
|
@lru_cache(maxsize=2)
|
710
1684
|
def xyxyxyxyn(self):
|
711
|
-
"""
|
1685
|
+
"""
|
1686
|
+
Converts rotated bounding boxes to normalized xyxyxyxy format.
|
1687
|
+
|
1688
|
+
Returns:
|
1689
|
+
(torch.Tensor | numpy.ndarray): Normalized rotated bounding boxes in xyxyxyxy format with shape (N, 4, 2),
|
1690
|
+
where N is the number of boxes. Each box is represented by 4 points (x, y), normalized relative to
|
1691
|
+
the original image dimensions.
|
1692
|
+
|
1693
|
+
Examples:
|
1694
|
+
>>> obb = OBB(torch.rand(10, 7), orig_shape=(640, 480)) # 10 random OBBs
|
1695
|
+
>>> normalized_boxes = obb.xyxyxyxyn
|
1696
|
+
>>> print(normalized_boxes.shape)
|
1697
|
+
torch.Size([10, 4, 2])
|
1698
|
+
"""
|
712
1699
|
xyxyxyxyn = self.xyxyxyxy.clone() if isinstance(self.xyxyxyxy, torch.Tensor) else np.copy(self.xyxyxyxy)
|
713
1700
|
xyxyxyxyn[..., 0] /= self.orig_shape[1]
|
714
1701
|
xyxyxyxyn[..., 1] /= self.orig_shape[0]
|
@@ -718,13 +1705,36 @@ class OBB(BaseTensor):
|
|
718
1705
|
@lru_cache(maxsize=2)
|
719
1706
|
def xyxy(self):
|
720
1707
|
"""
|
721
|
-
|
1708
|
+
Converts oriented bounding boxes (OBB) to axis-aligned bounding boxes in xyxy format.
|
1709
|
+
|
1710
|
+
This property calculates the minimal enclosing rectangle for each oriented bounding box and returns it in
|
1711
|
+
xyxy format (x1, y1, x2, y2). This is useful for operations that require axis-aligned bounding boxes, such
|
1712
|
+
as IoU calculation with non-rotated boxes.
|
722
1713
|
|
723
|
-
|
1714
|
+
Returns:
|
1715
|
+
(torch.Tensor | numpy.ndarray): Axis-aligned bounding boxes in xyxy format with shape (N, 4), where N
|
1716
|
+
is the number of boxes. Each row contains [x1, y1, x2, y2] coordinates.
|
1717
|
+
|
1718
|
+
Examples:
|
1719
|
+
>>> import torch
|
1720
|
+
>>> from ultralytics import YOLO
|
1721
|
+
>>> model = YOLO("yolov8n-obb.pt")
|
1722
|
+
>>> results = model("path/to/image.jpg")
|
1723
|
+
>>> for result in results:
|
1724
|
+
... obb = result.obb
|
1725
|
+
... if obb is not None:
|
1726
|
+
... xyxy_boxes = obb.xyxy
|
1727
|
+
... print(xyxy_boxes.shape) # (N, 4)
|
1728
|
+
|
1729
|
+
Notes:
|
1730
|
+
- This method approximates the OBB by its minimal enclosing rectangle.
|
1731
|
+
- The returned format is compatible with standard object detection metrics and visualization tools.
|
1732
|
+
- The property uses caching to improve performance for repeated access.
|
724
1733
|
"""
|
725
|
-
|
726
|
-
|
727
|
-
|
728
|
-
|
729
|
-
|
730
|
-
|
1734
|
+
x = self.xyxyxyxy[..., 0]
|
1735
|
+
y = self.xyxyxyxy[..., 1]
|
1736
|
+
return (
|
1737
|
+
torch.stack([x.amin(1), y.amin(1), x.amax(1), y.amax(1)], -1)
|
1738
|
+
if isinstance(x, torch.Tensor)
|
1739
|
+
else np.stack([x.min(1), y.min(1), x.max(1), y.max(1)], -1)
|
1740
|
+
)
|