ultralytics 8.1.28__py3-none-any.whl → 8.3.62__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.
Files changed (247) hide show
  1. tests/__init__.py +22 -0
  2. tests/conftest.py +83 -0
  3. tests/test_cli.py +122 -0
  4. tests/test_cuda.py +155 -0
  5. tests/test_engine.py +131 -0
  6. tests/test_exports.py +216 -0
  7. tests/test_integrations.py +150 -0
  8. tests/test_python.py +615 -0
  9. tests/test_solutions.py +94 -0
  10. ultralytics/__init__.py +11 -8
  11. ultralytics/cfg/__init__.py +569 -131
  12. ultralytics/cfg/datasets/Argoverse.yaml +2 -1
  13. ultralytics/cfg/datasets/DOTAv1.5.yaml +3 -2
  14. ultralytics/cfg/datasets/DOTAv1.yaml +3 -2
  15. ultralytics/cfg/datasets/GlobalWheat2020.yaml +3 -2
  16. ultralytics/cfg/datasets/ImageNet.yaml +2 -1
  17. ultralytics/cfg/datasets/Objects365.yaml +5 -4
  18. ultralytics/cfg/datasets/SKU-110K.yaml +2 -1
  19. ultralytics/cfg/datasets/VOC.yaml +3 -2
  20. ultralytics/cfg/datasets/VisDrone.yaml +6 -5
  21. ultralytics/cfg/datasets/african-wildlife.yaml +25 -0
  22. ultralytics/cfg/datasets/brain-tumor.yaml +23 -0
  23. ultralytics/cfg/datasets/carparts-seg.yaml +3 -2
  24. ultralytics/cfg/datasets/coco-pose.yaml +7 -6
  25. ultralytics/cfg/datasets/coco.yaml +3 -2
  26. ultralytics/cfg/datasets/coco128-seg.yaml +4 -3
  27. ultralytics/cfg/datasets/coco128.yaml +4 -3
  28. ultralytics/cfg/datasets/coco8-pose.yaml +3 -2
  29. ultralytics/cfg/datasets/coco8-seg.yaml +3 -2
  30. ultralytics/cfg/datasets/coco8.yaml +3 -2
  31. ultralytics/cfg/datasets/crack-seg.yaml +3 -2
  32. ultralytics/cfg/datasets/dog-pose.yaml +24 -0
  33. ultralytics/cfg/datasets/dota8.yaml +3 -2
  34. ultralytics/cfg/datasets/hand-keypoints.yaml +26 -0
  35. ultralytics/cfg/datasets/lvis.yaml +1236 -0
  36. ultralytics/cfg/datasets/medical-pills.yaml +22 -0
  37. ultralytics/cfg/datasets/open-images-v7.yaml +2 -1
  38. ultralytics/cfg/datasets/package-seg.yaml +5 -4
  39. ultralytics/cfg/datasets/signature.yaml +21 -0
  40. ultralytics/cfg/datasets/tiger-pose.yaml +3 -2
  41. ultralytics/cfg/datasets/xView.yaml +2 -1
  42. ultralytics/cfg/default.yaml +14 -11
  43. ultralytics/cfg/models/11/yolo11-cls-resnet18.yaml +24 -0
  44. ultralytics/cfg/models/11/yolo11-cls.yaml +33 -0
  45. ultralytics/cfg/models/11/yolo11-obb.yaml +50 -0
  46. ultralytics/cfg/models/11/yolo11-pose.yaml +51 -0
  47. ultralytics/cfg/models/11/yolo11-seg.yaml +50 -0
  48. ultralytics/cfg/models/11/yolo11.yaml +50 -0
  49. ultralytics/cfg/models/rt-detr/rtdetr-l.yaml +5 -2
  50. ultralytics/cfg/models/rt-detr/rtdetr-resnet101.yaml +5 -2
  51. ultralytics/cfg/models/rt-detr/rtdetr-resnet50.yaml +5 -2
  52. ultralytics/cfg/models/rt-detr/rtdetr-x.yaml +5 -2
  53. ultralytics/cfg/models/v10/yolov10b.yaml +45 -0
  54. ultralytics/cfg/models/v10/yolov10l.yaml +45 -0
  55. ultralytics/cfg/models/v10/yolov10m.yaml +45 -0
  56. ultralytics/cfg/models/v10/yolov10n.yaml +45 -0
  57. ultralytics/cfg/models/v10/yolov10s.yaml +45 -0
  58. ultralytics/cfg/models/v10/yolov10x.yaml +45 -0
  59. ultralytics/cfg/models/v3/yolov3-spp.yaml +5 -2
  60. ultralytics/cfg/models/v3/yolov3-tiny.yaml +5 -2
  61. ultralytics/cfg/models/v3/yolov3.yaml +5 -2
  62. ultralytics/cfg/models/v5/yolov5-p6.yaml +5 -2
  63. ultralytics/cfg/models/v5/yolov5.yaml +5 -2
  64. ultralytics/cfg/models/v6/yolov6.yaml +5 -2
  65. ultralytics/cfg/models/v8/yolov8-cls-resnet101.yaml +5 -2
  66. ultralytics/cfg/models/v8/yolov8-cls-resnet50.yaml +5 -2
  67. ultralytics/cfg/models/v8/yolov8-cls.yaml +5 -2
  68. ultralytics/cfg/models/v8/yolov8-ghost-p2.yaml +6 -2
  69. ultralytics/cfg/models/v8/yolov8-ghost-p6.yaml +6 -2
  70. ultralytics/cfg/models/v8/yolov8-ghost.yaml +5 -2
  71. ultralytics/cfg/models/v8/yolov8-obb.yaml +5 -2
  72. ultralytics/cfg/models/v8/yolov8-p2.yaml +5 -2
  73. ultralytics/cfg/models/v8/yolov8-p6.yaml +10 -7
  74. ultralytics/cfg/models/v8/yolov8-pose-p6.yaml +5 -2
  75. ultralytics/cfg/models/v8/yolov8-pose.yaml +5 -2
  76. ultralytics/cfg/models/v8/yolov8-rtdetr.yaml +5 -2
  77. ultralytics/cfg/models/v8/yolov8-seg-p6.yaml +5 -2
  78. ultralytics/cfg/models/v8/yolov8-seg.yaml +5 -2
  79. ultralytics/cfg/models/v8/yolov8-world.yaml +5 -2
  80. ultralytics/cfg/models/v8/yolov8-worldv2.yaml +5 -2
  81. ultralytics/cfg/models/v8/yolov8.yaml +5 -2
  82. ultralytics/cfg/models/v9/yolov9c-seg.yaml +41 -0
  83. ultralytics/cfg/models/v9/yolov9c.yaml +30 -25
  84. ultralytics/cfg/models/v9/yolov9e-seg.yaml +64 -0
  85. ultralytics/cfg/models/v9/yolov9e.yaml +46 -42
  86. ultralytics/cfg/models/v9/yolov9m.yaml +41 -0
  87. ultralytics/cfg/models/v9/yolov9s.yaml +41 -0
  88. ultralytics/cfg/models/v9/yolov9t.yaml +41 -0
  89. ultralytics/cfg/solutions/default.yaml +24 -0
  90. ultralytics/cfg/trackers/botsort.yaml +8 -5
  91. ultralytics/cfg/trackers/bytetrack.yaml +8 -5
  92. ultralytics/data/__init__.py +14 -3
  93. ultralytics/data/annotator.py +37 -15
  94. ultralytics/data/augment.py +1783 -289
  95. ultralytics/data/base.py +62 -27
  96. ultralytics/data/build.py +36 -8
  97. ultralytics/data/converter.py +196 -36
  98. ultralytics/data/dataset.py +233 -94
  99. ultralytics/data/loaders.py +199 -96
  100. ultralytics/data/split_dota.py +39 -29
  101. ultralytics/data/utils.py +110 -40
  102. ultralytics/engine/__init__.py +1 -1
  103. ultralytics/engine/exporter.py +569 -242
  104. ultralytics/engine/model.py +604 -252
  105. ultralytics/engine/predictor.py +22 -11
  106. ultralytics/engine/results.py +1228 -218
  107. ultralytics/engine/trainer.py +190 -129
  108. ultralytics/engine/tuner.py +18 -18
  109. ultralytics/engine/validator.py +18 -15
  110. ultralytics/hub/__init__.py +31 -13
  111. ultralytics/hub/auth.py +11 -7
  112. ultralytics/hub/google/__init__.py +159 -0
  113. ultralytics/hub/session.py +128 -94
  114. ultralytics/hub/utils.py +20 -21
  115. ultralytics/models/__init__.py +4 -2
  116. ultralytics/models/fastsam/__init__.py +2 -3
  117. ultralytics/models/fastsam/model.py +26 -4
  118. ultralytics/models/fastsam/predict.py +127 -63
  119. ultralytics/models/fastsam/utils.py +1 -44
  120. ultralytics/models/fastsam/val.py +1 -1
  121. ultralytics/models/nas/__init__.py +1 -1
  122. ultralytics/models/nas/model.py +21 -10
  123. ultralytics/models/nas/predict.py +3 -6
  124. ultralytics/models/nas/val.py +4 -4
  125. ultralytics/models/rtdetr/__init__.py +1 -1
  126. ultralytics/models/rtdetr/model.py +1 -1
  127. ultralytics/models/rtdetr/predict.py +6 -8
  128. ultralytics/models/rtdetr/train.py +6 -2
  129. ultralytics/models/rtdetr/val.py +3 -3
  130. ultralytics/models/sam/__init__.py +3 -3
  131. ultralytics/models/sam/amg.py +29 -23
  132. ultralytics/models/sam/build.py +211 -13
  133. ultralytics/models/sam/model.py +91 -30
  134. ultralytics/models/sam/modules/__init__.py +1 -1
  135. ultralytics/models/sam/modules/blocks.py +1129 -0
  136. ultralytics/models/sam/modules/decoders.py +381 -53
  137. ultralytics/models/sam/modules/encoders.py +515 -324
  138. ultralytics/models/sam/modules/memory_attention.py +237 -0
  139. ultralytics/models/sam/modules/sam.py +969 -21
  140. ultralytics/models/sam/modules/tiny_encoder.py +425 -154
  141. ultralytics/models/sam/modules/transformer.py +159 -60
  142. ultralytics/models/sam/modules/utils.py +293 -0
  143. ultralytics/models/sam/predict.py +1263 -132
  144. ultralytics/models/utils/__init__.py +1 -1
  145. ultralytics/models/utils/loss.py +36 -24
  146. ultralytics/models/utils/ops.py +3 -7
  147. ultralytics/models/yolo/__init__.py +3 -3
  148. ultralytics/models/yolo/classify/__init__.py +1 -1
  149. ultralytics/models/yolo/classify/predict.py +7 -8
  150. ultralytics/models/yolo/classify/train.py +17 -22
  151. ultralytics/models/yolo/classify/val.py +8 -4
  152. ultralytics/models/yolo/detect/__init__.py +1 -1
  153. ultralytics/models/yolo/detect/predict.py +3 -5
  154. ultralytics/models/yolo/detect/train.py +11 -4
  155. ultralytics/models/yolo/detect/val.py +90 -52
  156. ultralytics/models/yolo/model.py +14 -9
  157. ultralytics/models/yolo/obb/__init__.py +1 -1
  158. ultralytics/models/yolo/obb/predict.py +2 -2
  159. ultralytics/models/yolo/obb/train.py +5 -3
  160. ultralytics/models/yolo/obb/val.py +41 -23
  161. ultralytics/models/yolo/pose/__init__.py +1 -1
  162. ultralytics/models/yolo/pose/predict.py +3 -5
  163. ultralytics/models/yolo/pose/train.py +2 -2
  164. ultralytics/models/yolo/pose/val.py +51 -17
  165. ultralytics/models/yolo/segment/__init__.py +1 -1
  166. ultralytics/models/yolo/segment/predict.py +3 -5
  167. ultralytics/models/yolo/segment/train.py +2 -2
  168. ultralytics/models/yolo/segment/val.py +60 -19
  169. ultralytics/models/yolo/world/__init__.py +5 -0
  170. ultralytics/models/yolo/world/train.py +92 -0
  171. ultralytics/models/yolo/world/train_world.py +109 -0
  172. ultralytics/nn/__init__.py +1 -1
  173. ultralytics/nn/autobackend.py +228 -93
  174. ultralytics/nn/modules/__init__.py +39 -14
  175. ultralytics/nn/modules/activation.py +21 -0
  176. ultralytics/nn/modules/block.py +527 -67
  177. ultralytics/nn/modules/conv.py +24 -7
  178. ultralytics/nn/modules/head.py +177 -34
  179. ultralytics/nn/modules/transformer.py +6 -5
  180. ultralytics/nn/modules/utils.py +1 -2
  181. ultralytics/nn/tasks.py +225 -77
  182. ultralytics/solutions/__init__.py +30 -1
  183. ultralytics/solutions/ai_gym.py +96 -143
  184. ultralytics/solutions/analytics.py +247 -0
  185. ultralytics/solutions/distance_calculation.py +78 -135
  186. ultralytics/solutions/heatmap.py +93 -247
  187. ultralytics/solutions/object_counter.py +184 -259
  188. ultralytics/solutions/parking_management.py +246 -0
  189. ultralytics/solutions/queue_management.py +112 -0
  190. ultralytics/solutions/region_counter.py +116 -0
  191. ultralytics/solutions/security_alarm.py +144 -0
  192. ultralytics/solutions/solutions.py +178 -0
  193. ultralytics/solutions/speed_estimation.py +86 -174
  194. ultralytics/solutions/streamlit_inference.py +190 -0
  195. ultralytics/solutions/trackzone.py +68 -0
  196. ultralytics/trackers/__init__.py +1 -1
  197. ultralytics/trackers/basetrack.py +32 -13
  198. ultralytics/trackers/bot_sort.py +61 -28
  199. ultralytics/trackers/byte_tracker.py +83 -51
  200. ultralytics/trackers/track.py +21 -6
  201. ultralytics/trackers/utils/__init__.py +1 -1
  202. ultralytics/trackers/utils/gmc.py +62 -48
  203. ultralytics/trackers/utils/kalman_filter.py +166 -35
  204. ultralytics/trackers/utils/matching.py +40 -21
  205. ultralytics/utils/__init__.py +511 -239
  206. ultralytics/utils/autobatch.py +40 -22
  207. ultralytics/utils/benchmarks.py +266 -85
  208. ultralytics/utils/callbacks/__init__.py +1 -1
  209. ultralytics/utils/callbacks/base.py +1 -3
  210. ultralytics/utils/callbacks/clearml.py +7 -6
  211. ultralytics/utils/callbacks/comet.py +39 -17
  212. ultralytics/utils/callbacks/dvc.py +1 -1
  213. ultralytics/utils/callbacks/hub.py +16 -16
  214. ultralytics/utils/callbacks/mlflow.py +28 -24
  215. ultralytics/utils/callbacks/neptune.py +6 -2
  216. ultralytics/utils/callbacks/raytune.py +3 -4
  217. ultralytics/utils/callbacks/tensorboard.py +18 -18
  218. ultralytics/utils/callbacks/wb.py +27 -20
  219. ultralytics/utils/checks.py +160 -100
  220. ultralytics/utils/dist.py +2 -1
  221. ultralytics/utils/downloads.py +44 -37
  222. ultralytics/utils/errors.py +1 -1
  223. ultralytics/utils/files.py +72 -38
  224. ultralytics/utils/instance.py +41 -19
  225. ultralytics/utils/loss.py +84 -56
  226. ultralytics/utils/metrics.py +61 -56
  227. ultralytics/utils/ops.py +94 -89
  228. ultralytics/utils/patches.py +30 -14
  229. ultralytics/utils/plotting.py +600 -269
  230. ultralytics/utils/tal.py +67 -26
  231. ultralytics/utils/torch_utils.py +302 -102
  232. ultralytics/utils/triton.py +2 -1
  233. ultralytics/utils/tuner.py +21 -12
  234. ultralytics-8.3.62.dist-info/METADATA +370 -0
  235. ultralytics-8.3.62.dist-info/RECORD +241 -0
  236. {ultralytics-8.1.28.dist-info → ultralytics-8.3.62.dist-info}/WHEEL +1 -1
  237. ultralytics/data/explorer/__init__.py +0 -5
  238. ultralytics/data/explorer/explorer.py +0 -472
  239. ultralytics/data/explorer/gui/__init__.py +0 -1
  240. ultralytics/data/explorer/gui/dash.py +0 -268
  241. ultralytics/data/explorer/utils.py +0 -166
  242. ultralytics/models/fastsam/prompt.py +0 -357
  243. ultralytics-8.1.28.dist-info/METADATA +0 -373
  244. ultralytics-8.1.28.dist-info/RECORD +0 -197
  245. {ultralytics-8.1.28.dist-info → ultralytics-8.3.62.dist-info}/LICENSE +0 -0
  246. {ultralytics-8.1.28.dist-info → ultralytics-8.3.62.dist-info}/entry_points.txt +0 -0
  247. {ultralytics-8.1.28.dist-info → ultralytics-8.3.62.dist-info}/top_level.txt +0 -0
@@ -1,4 +1,4 @@
1
- # Ultralytics YOLO 🚀, AGPL-3.0 license
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
- """Base tensor class with additional methods for easy manipulation and device handling."""
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): Predictions, such as bboxes, masks and keypoints.
30
- orig_shape (tuple): Original shape of image.
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
- """Return the shape of the data tensor."""
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
- """Return a copy of the tensor on CPU memory."""
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
- """Return a copy of the tensor as a numpy array."""
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
- """Return a copy of the tensor on GPU memory."""
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
- """Return a copy of the tensor with the specified device and dtype."""
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
- """Return the length of the data tensor."""
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
- """Return a BaseTensor with the specified index of the data tensor."""
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 (tuple): Original image shape in (height, width) format.
73
- boxes (Boxes, optional): Object containing detection bounding boxes.
74
- masks (Masks, optional): Object containing detection masks.
75
- probs (Probs, optional): Object containing class probabilities for classification tasks.
76
- keypoints (Keypoints, optional): Object containing detected keypoints for each object.
77
- speed (dict): Dictionary of preprocess, inference, and postprocess speeds (ms/image).
78
- names (dict): Dictionary of class names.
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(boxes=None, masks=None, probs=None, obb=None): Updates object attributes with new detection results.
83
- cpu(): Returns a copy of the Results object with all tensors on CPU memory.
84
- numpy(): Returns a copy of the Results object with all tensors as numpy arrays.
85
- cuda(): Returns a copy of the Results object with all tensors on GPU memory.
86
- to(*args, **kwargs): Returns a copy of the Results object with tensors on a specified device and dtype.
87
- new(): Returns a new Results object with the same image, path, and names.
88
- plot(...): Plots detection results on an input image, returning an annotated image.
89
- show(): Show annotated results to screen.
90
- save(filename): Save annotated results to file.
91
- verbose(): Returns a log string for each task, detailing detections and classifications.
92
- save_txt(txt_file, save_conf=False): Saves detection results to a text file.
93
- save_crop(save_dir, file_name=Path("im.jpg")): Saves cropped detection images.
94
- tojson(normalize=False): Converts detection results to JSON format.
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__(self, orig_img, path, names, boxes=None, masks=None, probs=None, keypoints=None, obb=None) -> None:
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 (dict): A dictionary of class names.
105
- boxes (torch.tensor, optional): A 2D tensor of bounding box coordinates for each detection.
106
- masks (torch.tensor, optional): A 3D tensor of detection masks, where each mask is a binary image.
107
- probs (torch.tensor, optional): A 1D tensor of probabilities of each class for classification task.
108
- keypoints (torch.tensor, optional): A 2D tensor of keypoint coordinates for each detection.
109
- obb (torch.tensor, optional): A 2D tensor of oriented bounding box coordinates for each detection.
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} # milliseconds per image
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
- """Return a Results object for the specified index."""
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
- """Return the number of detections in the Results object."""
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
- """Update the boxes, masks, and probs attributes of the Results object."""
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. This
149
- function is internally called by methods like .to(), .cuda(), .cpu(), etc.
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
- """Return a copy of the Results object with all tensors on CPU memory."""
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
- """Return a copy of the Results object with all tensors as numpy arrays."""
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
- """Return a copy of the Results object with all tensors on GPU memory."""
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
- """Return a copy of the Results object with tensors on the specified device and dtype."""
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
- """Return a new Results object with the same image, path, and names."""
184
- return Results(orig_img=self.orig_img, path=self.path, names=self.names)
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 the detection results on an input RGB image. Accepts a numpy array (cv2) or a PIL Image.
467
+ Plots detection results on an input RGB image.
207
468
 
208
469
  Args:
209
- conf (bool): Whether to plot the detection confidence score.
210
- line_width (float, optional): The line width of the bounding boxes. If None, it is scaled to the image size.
211
- font_size (float, optional): The font size of the text. If None, it is scaled to the image size.
212
- font (str): The font to use for the text.
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 (numpy.ndarray): Plot to another image. if not, plot to original image.
215
- im_gpu (torch.Tensor): Normalized image in gpu with shape (1, 3, 640, 640), for faster mask plotting.
216
- kpt_radius (int, optional): Radius of the drawn keypoints. Default is 5.
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 the label of bounding boxes.
219
- boxes (bool): Whether to plot the bounding boxes.
220
- masks (bool): Whether to plot the masks.
221
- probs (bool): Whether to plot classification probability
222
- show (bool): Whether to display the annotated image directly.
223
- save (bool): Whether to save the annotated image to `filename`.
224
- filename (str): Filename to save image to if save is True.
225
-
226
- Returns:
227
- (numpy.ndarray): A numpy array of the annotated image.
228
-
229
- Example:
230
- ```python
231
- from PIL import Image
232
- from ultralytics import YOLO
233
-
234
- model = YOLO('yolov8n.pt')
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 = pred_boxes.cls if pred_boxes else range(len(pred_masks))
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, conf, id = int(d.cls), float(d.conf) if conf else None, None if d.id is None else int(d.id.item())
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} {conf:.2f}" if conf else name) if labels else None
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(box, label, color=colors(c, True), rotated=is_obb)
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(k, self.orig_shape, radius=kpt_radius, kpt_line=kpt_line)
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
- """Show annotated results image."""
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
- """Save annotated results image."""
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
- """Return log string for each task."""
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 predictions into txt file.
667
+ Save detection results to a text file.
333
668
 
334
669
  Args:
335
- txt_file (str): txt file path.
336
- save_conf (bool): save confidence score or not.
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
- Save cropped predictions to `save_dir/cls/file_name.jpg`.
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 | pathlib.Path): Save path.
372
- file_name (str | pathlib.Path): File name.
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)] / f"{Path(file_name)}.jpg",
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 tojson(self, normalize=False):
389
- """Convert the object to JSON format."""
390
- if self.probs is not None:
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
- import json
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
- data = self.boxes.data.cpu().tolist()
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
- box = {"x1": row[0] / w, "y1": row[1] / h, "x2": row[2] / w, "y2": row[3] / h}
402
- conf = row[-2]
403
- class_id = int(row[-1])
404
- name = self.names[class_id]
405
- result = {"name": name, "class": class_id, "confidence": conf, "box": box}
406
- if self.boxes.is_track:
407
- result["track_id"] = int(row[-3]) # track ID
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
- x, y = self.masks.xy[i][:, 0], self.masks.xy[i][:, 1] # numpy array
410
- result["segments"] = {"x": (x / w).tolist(), "y": (y / h).tolist()}
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"] = {"x": (x / w).tolist(), "y": (y / h).tolist(), "visible": visible.tolist()}
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
- # Convert detections to JSON
417
- return json.dumps(results, indent=2)
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
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
- Manages detection boxes, providing easy access and manipulation of box coordinates, confidence scores, class
423
- identifiers, and optional tracking IDs. Supports multiple formats for box coordinates, including both absolute and
424
- normalized forms.
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 their associated data.
428
- orig_shape (tuple): The original image size as a tuple (height, width), used for normalization.
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, optional): Tracking IDs for each box, if available.
436
- xywh (torch.Tensor | numpy.ndarray): Boxes in [x, y, width, height] format, calculated on demand.
437
- xyxyn (torch.Tensor | numpy.ndarray): Normalized [x1, y1, x2, y2] boxes, relative to `orig_shape`.
438
- xywhn (torch.Tensor | numpy.ndarray): Normalized [x, y, width, height] boxes, relative to `orig_shape`.
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(): Moves the boxes to CPU memory.
442
- numpy(): Converts the boxes to a numpy array format.
443
- cuda(): Moves the boxes to CUDA (GPU) memory.
444
- to(device, dtype=None): Moves the boxes to the specified device.
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 | numpy.ndarray): A tensor or numpy array containing the detection boxes, with
453
- shape (num_boxes, 6) or (num_boxes, 7). The last two columns contain confidence and class values.
454
- If present, the third last column contains track IDs.
455
- orig_shape (tuple): Original image size, in the format (height, width).
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 (6, 7), f"expected 6 or 7 values but got {n}" # xyxy, track_id, conf, cls
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
- """Return the boxes in xyxy format."""
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
- """Return the confidence values of the boxes."""
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
- """Return the class values of the boxes."""
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
- """Return the track IDs of the boxes (if available)."""
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
- """Return the boxes in xywh format."""
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
- """Return the boxes in xyxy format normalized by original image size."""
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
- """Return the boxes in xywh format normalized by original image size."""
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
- xy (list): A list of segments in pixel coordinates.
516
- xyn (list): A list of normalized segments.
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 masks tensor on CPU memory.
520
- numpy(): Returns the masks tensor as a numpy array.
521
- cuda(): Returns the masks tensor on GPU memory.
522
- to(device, dtype): Returns the masks tensor with the specified device and dtype.
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
- """Initialize the Masks class with the given masks tensor and original image shape."""
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
- """Return normalized segments."""
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
- """Return segments in pixel coordinates."""
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
- xy (torch.Tensor): A collection of keypoints containing x, y coordinates for each detection.
556
- xyn (torch.Tensor): A normalized version of xy with coordinates in the range [0, 1].
557
- conf (torch.Tensor): Confidence values associated with keypoints if available, otherwise None.
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(device, dtype): Returns a copy of the keypoints tensor with the specified device and dtype.
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
- """Initializes the Keypoints object with detection keypoints and original image size."""
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
- """Returns x, y coordinates of keypoints."""
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
- """Returns normalized x, y coordinates of keypoints."""
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
- """Returns confidence values of keypoints if available, else None."""
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 predictions.
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
- top1 (int): Index of the top 1 class.
605
- top5 (list[int]): Indices of the top 5 classes.
606
- top1conf (torch.Tensor): Confidence of the top 1 class.
607
- top5conf (torch.Tensor): Confidences of the top 5 classes.
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 probs tensor on CPU memory.
611
- numpy(): Returns a copy of the probs tensor as a numpy array.
612
- cuda(): Returns a copy of the probs tensor on GPU memory.
613
- to(): Returns a copy of the probs tensor with the specified device and dtype.
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
- """Initialize the Probs class with classification probabilities and optional original shape of the image."""
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
- """Return the index of top 1."""
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
- """Return the indices of top 5."""
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
- """Return the confidence of top 1."""
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
- """Return the confidences of top 5."""
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
- Args:
650
- boxes (torch.Tensor | numpy.ndarray): A tensor or numpy array containing the detection boxes,
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
- xywhr (torch.Tensor | numpy.ndarray): The boxes in [x_center, y_center, width, height, rotation] format.
657
- conf (torch.Tensor | numpy.ndarray): The confidence values of the boxes.
658
- cls (torch.Tensor | numpy.ndarray): The class values of the boxes.
659
- id (torch.Tensor | numpy.ndarray): The track IDs of the boxes (if available).
660
- xyxyxyxyn (torch.Tensor | numpy.ndarray): The rotated boxes in xyxyxyxy format normalized by orig image size.
661
- xyxyxyxy (torch.Tensor | numpy.ndarray): The rotated boxes in xyxyxyxy format.
662
- xyxy (torch.Tensor | numpy.ndarray): The horizontal boxes in xyxyxyxy format.
663
- data (torch.Tensor): The raw OBB tensor (alias for `boxes`).
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(): Move the object to CPU memory.
667
- numpy(): Convert the object to a numpy array.
668
- cuda(): Move the object to CUDA memory.
669
- to(*args, **kwargs): Move the object to the specified device.
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
- """Initialize the Boxes class."""
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 (7, 8), f"expected 7 or 8 values but got {n}" # xywh, rotation, track_id, conf, cls
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
- """Return the rotated boxes in xywhr format."""
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
- """Return the confidence values of the boxes."""
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
- """Return the class values of the boxes."""
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
- """Return the track IDs of the boxes (if available)."""
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
- """Return the boxes in xyxyxyxy format, (N, 4, 2)."""
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
- """Return the boxes in xyxyxyxy format, (N, 4, 2)."""
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
- Return the horizontal boxes in xyxy format, (N, 4).
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
- Accepts both torch and numpy boxes.
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
- x1 = self.xyxyxyxy[..., 0].min(1).values
726
- x2 = self.xyxyxyxy[..., 0].max(1).values
727
- y1 = self.xyxyxyxy[..., 1].min(1).values
728
- y2 = self.xyxyxyxy[..., 1].max(1).values
729
- xyxy = [x1, y1, x2, y2]
730
- return np.stack(xyxy, axis=-1) if isinstance(self.data, np.ndarray) else torch.stack(xyxy, dim=-1)
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
+ )