bplusplus 0.1.0__py3-none-any.whl → 1.1.0__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.

Potentially problematic release.


This version of bplusplus might be problematic. Click here for more details.

Files changed (95) hide show
  1. bplusplus/__init__.py +5 -3
  2. bplusplus/{collect_images.py → collect.py} +3 -3
  3. bplusplus/prepare.py +573 -0
  4. bplusplus/train_validate.py +8 -64
  5. bplusplus/yolov5detect/__init__.py +1 -0
  6. bplusplus/yolov5detect/detect.py +444 -0
  7. bplusplus/yolov5detect/export.py +1530 -0
  8. bplusplus/yolov5detect/insect.yaml +8 -0
  9. bplusplus/yolov5detect/models/__init__.py +0 -0
  10. bplusplus/yolov5detect/models/common.py +1109 -0
  11. bplusplus/yolov5detect/models/experimental.py +130 -0
  12. bplusplus/yolov5detect/models/hub/anchors.yaml +56 -0
  13. bplusplus/yolov5detect/models/hub/yolov3-spp.yaml +52 -0
  14. bplusplus/yolov5detect/models/hub/yolov3-tiny.yaml +42 -0
  15. bplusplus/yolov5detect/models/hub/yolov3.yaml +52 -0
  16. bplusplus/yolov5detect/models/hub/yolov5-bifpn.yaml +49 -0
  17. bplusplus/yolov5detect/models/hub/yolov5-fpn.yaml +43 -0
  18. bplusplus/yolov5detect/models/hub/yolov5-p2.yaml +55 -0
  19. bplusplus/yolov5detect/models/hub/yolov5-p34.yaml +42 -0
  20. bplusplus/yolov5detect/models/hub/yolov5-p6.yaml +57 -0
  21. bplusplus/yolov5detect/models/hub/yolov5-p7.yaml +68 -0
  22. bplusplus/yolov5detect/models/hub/yolov5-panet.yaml +49 -0
  23. bplusplus/yolov5detect/models/hub/yolov5l6.yaml +61 -0
  24. bplusplus/yolov5detect/models/hub/yolov5m6.yaml +61 -0
  25. bplusplus/yolov5detect/models/hub/yolov5n6.yaml +61 -0
  26. bplusplus/yolov5detect/models/hub/yolov5s-LeakyReLU.yaml +50 -0
  27. bplusplus/yolov5detect/models/hub/yolov5s-ghost.yaml +49 -0
  28. bplusplus/yolov5detect/models/hub/yolov5s-transformer.yaml +49 -0
  29. bplusplus/yolov5detect/models/hub/yolov5s6.yaml +61 -0
  30. bplusplus/yolov5detect/models/hub/yolov5x6.yaml +61 -0
  31. bplusplus/yolov5detect/models/segment/yolov5l-seg.yaml +49 -0
  32. bplusplus/yolov5detect/models/segment/yolov5m-seg.yaml +49 -0
  33. bplusplus/yolov5detect/models/segment/yolov5n-seg.yaml +49 -0
  34. bplusplus/yolov5detect/models/segment/yolov5s-seg.yaml +49 -0
  35. bplusplus/yolov5detect/models/segment/yolov5x-seg.yaml +49 -0
  36. bplusplus/yolov5detect/models/tf.py +797 -0
  37. bplusplus/yolov5detect/models/yolo.py +495 -0
  38. bplusplus/yolov5detect/models/yolov5l.yaml +49 -0
  39. bplusplus/yolov5detect/models/yolov5m.yaml +49 -0
  40. bplusplus/yolov5detect/models/yolov5n.yaml +49 -0
  41. bplusplus/yolov5detect/models/yolov5s.yaml +49 -0
  42. bplusplus/yolov5detect/models/yolov5x.yaml +49 -0
  43. bplusplus/yolov5detect/utils/__init__.py +97 -0
  44. bplusplus/yolov5detect/utils/activations.py +134 -0
  45. bplusplus/yolov5detect/utils/augmentations.py +448 -0
  46. bplusplus/yolov5detect/utils/autoanchor.py +175 -0
  47. bplusplus/yolov5detect/utils/autobatch.py +70 -0
  48. bplusplus/yolov5detect/utils/aws/__init__.py +0 -0
  49. bplusplus/yolov5detect/utils/aws/mime.sh +26 -0
  50. bplusplus/yolov5detect/utils/aws/resume.py +41 -0
  51. bplusplus/yolov5detect/utils/aws/userdata.sh +27 -0
  52. bplusplus/yolov5detect/utils/callbacks.py +72 -0
  53. bplusplus/yolov5detect/utils/dataloaders.py +1385 -0
  54. bplusplus/yolov5detect/utils/docker/Dockerfile +73 -0
  55. bplusplus/yolov5detect/utils/docker/Dockerfile-arm64 +40 -0
  56. bplusplus/yolov5detect/utils/docker/Dockerfile-cpu +42 -0
  57. bplusplus/yolov5detect/utils/downloads.py +136 -0
  58. bplusplus/yolov5detect/utils/flask_rest_api/README.md +70 -0
  59. bplusplus/yolov5detect/utils/flask_rest_api/example_request.py +17 -0
  60. bplusplus/yolov5detect/utils/flask_rest_api/restapi.py +49 -0
  61. bplusplus/yolov5detect/utils/general.py +1294 -0
  62. bplusplus/yolov5detect/utils/google_app_engine/Dockerfile +25 -0
  63. bplusplus/yolov5detect/utils/google_app_engine/additional_requirements.txt +6 -0
  64. bplusplus/yolov5detect/utils/google_app_engine/app.yaml +16 -0
  65. bplusplus/yolov5detect/utils/loggers/__init__.py +476 -0
  66. bplusplus/yolov5detect/utils/loggers/clearml/README.md +222 -0
  67. bplusplus/yolov5detect/utils/loggers/clearml/__init__.py +0 -0
  68. bplusplus/yolov5detect/utils/loggers/clearml/clearml_utils.py +230 -0
  69. bplusplus/yolov5detect/utils/loggers/clearml/hpo.py +90 -0
  70. bplusplus/yolov5detect/utils/loggers/comet/README.md +250 -0
  71. bplusplus/yolov5detect/utils/loggers/comet/__init__.py +551 -0
  72. bplusplus/yolov5detect/utils/loggers/comet/comet_utils.py +151 -0
  73. bplusplus/yolov5detect/utils/loggers/comet/hpo.py +126 -0
  74. bplusplus/yolov5detect/utils/loggers/comet/optimizer_config.json +135 -0
  75. bplusplus/yolov5detect/utils/loggers/wandb/__init__.py +0 -0
  76. bplusplus/yolov5detect/utils/loggers/wandb/wandb_utils.py +210 -0
  77. bplusplus/yolov5detect/utils/loss.py +259 -0
  78. bplusplus/yolov5detect/utils/metrics.py +381 -0
  79. bplusplus/yolov5detect/utils/plots.py +517 -0
  80. bplusplus/yolov5detect/utils/segment/__init__.py +0 -0
  81. bplusplus/yolov5detect/utils/segment/augmentations.py +100 -0
  82. bplusplus/yolov5detect/utils/segment/dataloaders.py +366 -0
  83. bplusplus/yolov5detect/utils/segment/general.py +160 -0
  84. bplusplus/yolov5detect/utils/segment/loss.py +198 -0
  85. bplusplus/yolov5detect/utils/segment/metrics.py +225 -0
  86. bplusplus/yolov5detect/utils/segment/plots.py +152 -0
  87. bplusplus/yolov5detect/utils/torch_utils.py +482 -0
  88. bplusplus/yolov5detect/utils/triton.py +90 -0
  89. bplusplus-1.1.0.dist-info/METADATA +179 -0
  90. bplusplus-1.1.0.dist-info/RECORD +92 -0
  91. bplusplus/build_model.py +0 -38
  92. bplusplus-0.1.0.dist-info/METADATA +0 -91
  93. bplusplus-0.1.0.dist-info/RECORD +0 -8
  94. {bplusplus-0.1.0.dist-info → bplusplus-1.1.0.dist-info}/LICENSE +0 -0
  95. {bplusplus-0.1.0.dist-info → bplusplus-1.1.0.dist-info}/WHEEL +0 -0
@@ -0,0 +1,1109 @@
1
+ # Ultralytics YOLOv5 🚀, AGPL-3.0 license
2
+ """Common modules."""
3
+
4
+ import ast
5
+ import contextlib
6
+ import json
7
+ import math
8
+ import platform
9
+ import warnings
10
+ import zipfile
11
+ from collections import OrderedDict, namedtuple
12
+ from copy import copy
13
+ from pathlib import Path
14
+ from urllib.parse import urlparse
15
+
16
+ import cv2
17
+ import numpy as np
18
+ import pandas as pd
19
+ import requests
20
+ import torch
21
+ import torch.nn as nn
22
+ from PIL import Image
23
+ from torch.cuda import amp
24
+
25
+ # Import 'ultralytics' package or install if missing
26
+ try:
27
+ import ultralytics
28
+
29
+ assert hasattr(ultralytics, "__version__") # verify package is not directory
30
+ except (ImportError, AssertionError):
31
+ import os
32
+
33
+ os.system("pip install -U ultralytics")
34
+ import ultralytics
35
+
36
+ from ultralytics.utils.plotting import Annotator, colors, save_one_box
37
+
38
+ from utils import TryExcept
39
+ from utils.dataloaders import exif_transpose, letterbox
40
+ from utils.general import (
41
+ LOGGER,
42
+ ROOT,
43
+ Profile,
44
+ check_requirements,
45
+ check_suffix,
46
+ check_version,
47
+ colorstr,
48
+ increment_path,
49
+ is_jupyter,
50
+ make_divisible,
51
+ non_max_suppression,
52
+ scale_boxes,
53
+ xywh2xyxy,
54
+ xyxy2xywh,
55
+ yaml_load,
56
+ )
57
+ from utils.torch_utils import copy_attr, smart_inference_mode
58
+
59
+
60
+ def autopad(k, p=None, d=1):
61
+ """
62
+ Pads kernel to 'same' output shape, adjusting for optional dilation; returns padding size.
63
+
64
+ `k`: kernel, `p`: padding, `d`: dilation.
65
+ """
66
+ if d > 1:
67
+ k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
68
+ if p is None:
69
+ p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
70
+ return p
71
+
72
+
73
+ class Conv(nn.Module):
74
+ """Applies a convolution, batch normalization, and activation function to an input tensor in a neural network."""
75
+
76
+ default_act = nn.SiLU() # default activation
77
+
78
+ def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
79
+ """Initializes a standard convolution layer with optional batch normalization and activation."""
80
+ super().__init__()
81
+ self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
82
+ self.bn = nn.BatchNorm2d(c2)
83
+ self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
84
+
85
+ def forward(self, x):
86
+ """Applies a convolution followed by batch normalization and an activation function to the input tensor `x`."""
87
+ return self.act(self.bn(self.conv(x)))
88
+
89
+ def forward_fuse(self, x):
90
+ """Applies a fused convolution and activation function to the input tensor `x`."""
91
+ return self.act(self.conv(x))
92
+
93
+
94
+ class DWConv(Conv):
95
+ """Implements a depth-wise convolution layer with optional activation for efficient spatial filtering."""
96
+
97
+ def __init__(self, c1, c2, k=1, s=1, d=1, act=True):
98
+ """Initializes a depth-wise convolution layer with optional activation; args: input channels (c1), output
99
+ channels (c2), kernel size (k), stride (s), dilation (d), and activation flag (act).
100
+ """
101
+ super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act)
102
+
103
+
104
+ class DWConvTranspose2d(nn.ConvTranspose2d):
105
+ """A depth-wise transpose convolutional layer for upsampling in neural networks, particularly in YOLOv5 models."""
106
+
107
+ def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0):
108
+ """Initializes a depth-wise transpose convolutional layer for YOLOv5; args: input channels (c1), output channels
109
+ (c2), kernel size (k), stride (s), input padding (p1), output padding (p2).
110
+ """
111
+ super().__init__(c1, c2, k, s, p1, p2, groups=math.gcd(c1, c2))
112
+
113
+
114
+ class TransformerLayer(nn.Module):
115
+ """Transformer layer with multihead attention and linear layers, optimized by removing LayerNorm."""
116
+
117
+ def __init__(self, c, num_heads):
118
+ """
119
+ Initializes a transformer layer, sans LayerNorm for performance, with multihead attention and linear layers.
120
+
121
+ See as described in https://arxiv.org/abs/2010.11929.
122
+ """
123
+ super().__init__()
124
+ self.q = nn.Linear(c, c, bias=False)
125
+ self.k = nn.Linear(c, c, bias=False)
126
+ self.v = nn.Linear(c, c, bias=False)
127
+ self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
128
+ self.fc1 = nn.Linear(c, c, bias=False)
129
+ self.fc2 = nn.Linear(c, c, bias=False)
130
+
131
+ def forward(self, x):
132
+ """Performs forward pass using MultiheadAttention and two linear transformations with residual connections."""
133
+ x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
134
+ x = self.fc2(self.fc1(x)) + x
135
+ return x
136
+
137
+
138
+ class TransformerBlock(nn.Module):
139
+ """A Transformer block for vision tasks with convolution, position embeddings, and Transformer layers."""
140
+
141
+ def __init__(self, c1, c2, num_heads, num_layers):
142
+ """Initializes a Transformer block for vision tasks, adapting dimensions if necessary and stacking specified
143
+ layers.
144
+ """
145
+ super().__init__()
146
+ self.conv = None
147
+ if c1 != c2:
148
+ self.conv = Conv(c1, c2)
149
+ self.linear = nn.Linear(c2, c2) # learnable position embedding
150
+ self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
151
+ self.c2 = c2
152
+
153
+ def forward(self, x):
154
+ """Processes input through an optional convolution, followed by Transformer layers and position embeddings for
155
+ object detection.
156
+ """
157
+ if self.conv is not None:
158
+ x = self.conv(x)
159
+ b, _, w, h = x.shape
160
+ p = x.flatten(2).permute(2, 0, 1)
161
+ return self.tr(p + self.linear(p)).permute(1, 2, 0).reshape(b, self.c2, w, h)
162
+
163
+
164
+ class Bottleneck(nn.Module):
165
+ """A bottleneck layer with optional shortcut and group convolution for efficient feature extraction."""
166
+
167
+ def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):
168
+ """Initializes a standard bottleneck layer with optional shortcut and group convolution, supporting channel
169
+ expansion.
170
+ """
171
+ super().__init__()
172
+ c_ = int(c2 * e) # hidden channels
173
+ self.cv1 = Conv(c1, c_, 1, 1)
174
+ self.cv2 = Conv(c_, c2, 3, 1, g=g)
175
+ self.add = shortcut and c1 == c2
176
+
177
+ def forward(self, x):
178
+ """Processes input through two convolutions, optionally adds shortcut if channel dimensions match; input is a
179
+ tensor.
180
+ """
181
+ return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
182
+
183
+
184
+ class BottleneckCSP(nn.Module):
185
+ """CSP bottleneck layer for feature extraction with cross-stage partial connections and optional shortcuts."""
186
+
187
+ def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
188
+ """Initializes CSP bottleneck with optional shortcuts; args: ch_in, ch_out, number of repeats, shortcut bool,
189
+ groups, expansion.
190
+ """
191
+ super().__init__()
192
+ c_ = int(c2 * e) # hidden channels
193
+ self.cv1 = Conv(c1, c_, 1, 1)
194
+ self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
195
+ self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
196
+ self.cv4 = Conv(2 * c_, c2, 1, 1)
197
+ self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
198
+ self.act = nn.SiLU()
199
+ self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
200
+
201
+ def forward(self, x):
202
+ """Performs forward pass by applying layers, activation, and concatenation on input x, returning feature-
203
+ enhanced output.
204
+ """
205
+ y1 = self.cv3(self.m(self.cv1(x)))
206
+ y2 = self.cv2(x)
207
+ return self.cv4(self.act(self.bn(torch.cat((y1, y2), 1))))
208
+
209
+
210
+ class CrossConv(nn.Module):
211
+ """Implements a cross convolution layer with downsampling, expansion, and optional shortcut."""
212
+
213
+ def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
214
+ """
215
+ Initializes CrossConv with downsampling, expanding, and optionally shortcutting; `c1` input, `c2` output
216
+ channels.
217
+
218
+ Inputs are ch_in, ch_out, kernel, stride, groups, expansion, shortcut.
219
+ """
220
+ super().__init__()
221
+ c_ = int(c2 * e) # hidden channels
222
+ self.cv1 = Conv(c1, c_, (1, k), (1, s))
223
+ self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
224
+ self.add = shortcut and c1 == c2
225
+
226
+ def forward(self, x):
227
+ """Performs feature sampling, expanding, and applies shortcut if channels match; expects `x` input tensor."""
228
+ return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
229
+
230
+
231
+ class C3(nn.Module):
232
+ """Implements a CSP Bottleneck module with three convolutions for enhanced feature extraction in neural networks."""
233
+
234
+ def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
235
+ """Initializes C3 module with options for channel count, bottleneck repetition, shortcut usage, group
236
+ convolutions, and expansion.
237
+ """
238
+ super().__init__()
239
+ c_ = int(c2 * e) # hidden channels
240
+ self.cv1 = Conv(c1, c_, 1, 1)
241
+ self.cv2 = Conv(c1, c_, 1, 1)
242
+ self.cv3 = Conv(2 * c_, c2, 1) # optional act=FReLU(c2)
243
+ self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
244
+
245
+ def forward(self, x):
246
+ """Performs forward propagation using concatenated outputs from two convolutions and a Bottleneck sequence."""
247
+ return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
248
+
249
+
250
+ class C3x(C3):
251
+ """Extends the C3 module with cross-convolutions for enhanced feature extraction in neural networks."""
252
+
253
+ def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
254
+ """Initializes C3x module with cross-convolutions, extending C3 with customizable channel dimensions, groups,
255
+ and expansion.
256
+ """
257
+ super().__init__(c1, c2, n, shortcut, g, e)
258
+ c_ = int(c2 * e)
259
+ self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))
260
+
261
+
262
+ class C3TR(C3):
263
+ """C3 module with TransformerBlock for enhanced feature extraction in object detection models."""
264
+
265
+ def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
266
+ """Initializes C3 module with TransformerBlock for enhanced feature extraction, accepts channel sizes, shortcut
267
+ config, group, and expansion.
268
+ """
269
+ super().__init__(c1, c2, n, shortcut, g, e)
270
+ c_ = int(c2 * e)
271
+ self.m = TransformerBlock(c_, c_, 4, n)
272
+
273
+
274
+ class C3SPP(C3):
275
+ """Extends the C3 module with an SPP layer for enhanced spatial feature extraction and customizable channels."""
276
+
277
+ def __init__(self, c1, c2, k=(5, 9, 13), n=1, shortcut=True, g=1, e=0.5):
278
+ """Initializes a C3 module with SPP layer for advanced spatial feature extraction, given channel sizes, kernel
279
+ sizes, shortcut, group, and expansion ratio.
280
+ """
281
+ super().__init__(c1, c2, n, shortcut, g, e)
282
+ c_ = int(c2 * e)
283
+ self.m = SPP(c_, c_, k)
284
+
285
+
286
+ class C3Ghost(C3):
287
+ """Implements a C3 module with Ghost Bottlenecks for efficient feature extraction in YOLOv5."""
288
+
289
+ def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
290
+ """Initializes YOLOv5's C3 module with Ghost Bottlenecks for efficient feature extraction."""
291
+ super().__init__(c1, c2, n, shortcut, g, e)
292
+ c_ = int(c2 * e) # hidden channels
293
+ self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
294
+
295
+
296
+ class SPP(nn.Module):
297
+ """Implements Spatial Pyramid Pooling (SPP) for feature extraction, ref: https://arxiv.org/abs/1406.4729."""
298
+
299
+ def __init__(self, c1, c2, k=(5, 9, 13)):
300
+ """Initializes SPP layer with Spatial Pyramid Pooling, ref: https://arxiv.org/abs/1406.4729, args: c1 (input channels), c2 (output channels), k (kernel sizes)."""
301
+ super().__init__()
302
+ c_ = c1 // 2 # hidden channels
303
+ self.cv1 = Conv(c1, c_, 1, 1)
304
+ self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
305
+ self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
306
+
307
+ def forward(self, x):
308
+ """Applies convolution and max pooling layers to the input tensor `x`, concatenates results, and returns output
309
+ tensor.
310
+ """
311
+ x = self.cv1(x)
312
+ with warnings.catch_warnings():
313
+ warnings.simplefilter("ignore") # suppress torch 1.9.0 max_pool2d() warning
314
+ return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
315
+
316
+
317
+ class SPPF(nn.Module):
318
+ """Implements a fast Spatial Pyramid Pooling (SPPF) layer for efficient feature extraction in YOLOv5 models."""
319
+
320
+ def __init__(self, c1, c2, k=5):
321
+ """
322
+ Initializes YOLOv5 SPPF layer with given channels and kernel size for YOLOv5 model, combining convolution and
323
+ max pooling.
324
+
325
+ Equivalent to SPP(k=(5, 9, 13)).
326
+ """
327
+ super().__init__()
328
+ c_ = c1 // 2 # hidden channels
329
+ self.cv1 = Conv(c1, c_, 1, 1)
330
+ self.cv2 = Conv(c_ * 4, c2, 1, 1)
331
+ self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
332
+
333
+ def forward(self, x):
334
+ """Processes input through a series of convolutions and max pooling operations for feature extraction."""
335
+ x = self.cv1(x)
336
+ with warnings.catch_warnings():
337
+ warnings.simplefilter("ignore") # suppress torch 1.9.0 max_pool2d() warning
338
+ y1 = self.m(x)
339
+ y2 = self.m(y1)
340
+ return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
341
+
342
+
343
+ class Focus(nn.Module):
344
+ """Focuses spatial information into channel space using slicing and convolution for efficient feature extraction."""
345
+
346
+ def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
347
+ """Initializes Focus module to concentrate width-height info into channel space with configurable convolution
348
+ parameters.
349
+ """
350
+ super().__init__()
351
+ self.conv = Conv(c1 * 4, c2, k, s, p, g, act=act)
352
+ # self.contract = Contract(gain=2)
353
+
354
+ def forward(self, x):
355
+ """Processes input through Focus mechanism, reshaping (b,c,w,h) to (b,4c,w/2,h/2) then applies convolution."""
356
+ return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))
357
+ # return self.conv(self.contract(x))
358
+
359
+
360
+ class GhostConv(nn.Module):
361
+ """Implements Ghost Convolution for efficient feature extraction, see https://github.com/huawei-noah/ghostnet."""
362
+
363
+ def __init__(self, c1, c2, k=1, s=1, g=1, act=True):
364
+ """Initializes GhostConv with in/out channels, kernel size, stride, groups, and activation; halves out channels
365
+ for efficiency.
366
+ """
367
+ super().__init__()
368
+ c_ = c2 // 2 # hidden channels
369
+ self.cv1 = Conv(c1, c_, k, s, None, g, act=act)
370
+ self.cv2 = Conv(c_, c_, 5, 1, None, c_, act=act)
371
+
372
+ def forward(self, x):
373
+ """Performs forward pass, concatenating outputs of two convolutions on input `x`: shape (B,C,H,W)."""
374
+ y = self.cv1(x)
375
+ return torch.cat((y, self.cv2(y)), 1)
376
+
377
+
378
+ class GhostBottleneck(nn.Module):
379
+ """Efficient bottleneck layer using Ghost Convolutions, see https://github.com/huawei-noah/ghostnet."""
380
+
381
+ def __init__(self, c1, c2, k=3, s=1):
382
+ """Initializes GhostBottleneck with ch_in `c1`, ch_out `c2`, kernel size `k`, stride `s`; see https://github.com/huawei-noah/ghostnet."""
383
+ super().__init__()
384
+ c_ = c2 // 2
385
+ self.conv = nn.Sequential(
386
+ GhostConv(c1, c_, 1, 1), # pw
387
+ DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
388
+ GhostConv(c_, c2, 1, 1, act=False),
389
+ ) # pw-linear
390
+ self.shortcut = (
391
+ nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
392
+ )
393
+
394
+ def forward(self, x):
395
+ """Processes input through conv and shortcut layers, returning their summed output."""
396
+ return self.conv(x) + self.shortcut(x)
397
+
398
+
399
+ class Contract(nn.Module):
400
+ """Contracts spatial dimensions into channel dimensions for efficient processing in neural networks."""
401
+
402
+ def __init__(self, gain=2):
403
+ """Initializes a layer to contract spatial dimensions (width-height) into channels, e.g., input shape
404
+ (1,64,80,80) to (1,256,40,40).
405
+ """
406
+ super().__init__()
407
+ self.gain = gain
408
+
409
+ def forward(self, x):
410
+ """Processes input tensor to expand channel dimensions by contracting spatial dimensions, yielding output shape
411
+ `(b, c*s*s, h//s, w//s)`.
412
+ """
413
+ b, c, h, w = x.size() # assert (h / s == 0) and (W / s == 0), 'Indivisible gain'
414
+ s = self.gain
415
+ x = x.view(b, c, h // s, s, w // s, s) # x(1,64,40,2,40,2)
416
+ x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
417
+ return x.view(b, c * s * s, h // s, w // s) # x(1,256,40,40)
418
+
419
+
420
+ class Expand(nn.Module):
421
+ """Expands spatial dimensions by redistributing channels, e.g., from (1,64,80,80) to (1,16,160,160)."""
422
+
423
+ def __init__(self, gain=2):
424
+ """
425
+ Initializes the Expand module to increase spatial dimensions by redistributing channels, with an optional gain
426
+ factor.
427
+
428
+ Example: x(1,64,80,80) to x(1,16,160,160).
429
+ """
430
+ super().__init__()
431
+ self.gain = gain
432
+
433
+ def forward(self, x):
434
+ """Processes input tensor x to expand spatial dimensions by redistributing channels, requiring C / gain^2 ==
435
+ 0.
436
+ """
437
+ b, c, h, w = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'
438
+ s = self.gain
439
+ x = x.view(b, s, s, c // s**2, h, w) # x(1,2,2,16,80,80)
440
+ x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)
441
+ return x.view(b, c // s**2, h * s, w * s) # x(1,16,160,160)
442
+
443
+
444
+ class Concat(nn.Module):
445
+ """Concatenates tensors along a specified dimension for efficient tensor manipulation in neural networks."""
446
+
447
+ def __init__(self, dimension=1):
448
+ """Initializes a Concat module to concatenate tensors along a specified dimension."""
449
+ super().__init__()
450
+ self.d = dimension
451
+
452
+ def forward(self, x):
453
+ """Concatenates a list of tensors along a specified dimension; `x` is a list of tensors, `dimension` is an
454
+ int.
455
+ """
456
+ return torch.cat(x, self.d)
457
+
458
+
459
+ class DetectMultiBackend(nn.Module):
460
+ """YOLOv5 MultiBackend class for inference on various backends including PyTorch, ONNX, TensorRT, and more."""
461
+
462
+ def __init__(self, weights="yolov5s.pt", device=torch.device("cpu"), dnn=False, data=None, fp16=False, fuse=True):
463
+ """Initializes DetectMultiBackend with support for various inference backends, including PyTorch and ONNX."""
464
+ # PyTorch: weights = *.pt
465
+ # TorchScript: *.torchscript
466
+ # ONNX Runtime: *.onnx
467
+ # ONNX OpenCV DNN: *.onnx --dnn
468
+ # OpenVINO: *_openvino_model
469
+ # CoreML: *.mlpackage
470
+ # TensorRT: *.engine
471
+ # TensorFlow SavedModel: *_saved_model
472
+ # TensorFlow GraphDef: *.pb
473
+ # TensorFlow Lite: *.tflite
474
+ # TensorFlow Edge TPU: *_edgetpu.tflite
475
+ # PaddlePaddle: *_paddle_model
476
+ from models.experimental import attempt_download, attempt_load # scoped to avoid circular import
477
+
478
+ super().__init__()
479
+ w = str(weights[0] if isinstance(weights, list) else weights)
480
+ pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, triton = self._model_type(w)
481
+ fp16 &= pt or jit or onnx or engine or triton # FP16
482
+ nhwc = coreml or saved_model or pb or tflite or edgetpu # BHWC formats (vs torch BCWH)
483
+ stride = 32 # default stride
484
+ cuda = torch.cuda.is_available() and device.type != "cpu" # use CUDA
485
+ if not (pt or triton):
486
+ w = attempt_download(w) # download if not local
487
+
488
+ if pt: # PyTorch
489
+ model = attempt_load(weights if isinstance(weights, list) else w, device=device, inplace=True, fuse=fuse)
490
+ stride = max(int(model.stride.max()), 32) # model stride
491
+ names = model.module.names if hasattr(model, "module") else model.names # get class names
492
+ model.half() if fp16 else model.float()
493
+ self.model = model # explicitly assign for to(), cpu(), cuda(), half()
494
+ elif jit: # TorchScript
495
+ LOGGER.info(f"Loading {w} for TorchScript inference...")
496
+ extra_files = {"config.txt": ""} # model metadata
497
+ model = torch.jit.load(w, _extra_files=extra_files, map_location=device)
498
+ model.half() if fp16 else model.float()
499
+ if extra_files["config.txt"]: # load metadata dict
500
+ d = json.loads(
501
+ extra_files["config.txt"],
502
+ object_hook=lambda d: {int(k) if k.isdigit() else k: v for k, v in d.items()},
503
+ )
504
+ stride, names = int(d["stride"]), d["names"]
505
+ elif dnn: # ONNX OpenCV DNN
506
+ LOGGER.info(f"Loading {w} for ONNX OpenCV DNN inference...")
507
+ check_requirements("opencv-python>=4.5.4")
508
+ net = cv2.dnn.readNetFromONNX(w)
509
+ elif onnx: # ONNX Runtime
510
+ LOGGER.info(f"Loading {w} for ONNX Runtime inference...")
511
+ check_requirements(("onnx", "onnxruntime-gpu" if cuda else "onnxruntime"))
512
+ import onnxruntime
513
+
514
+ providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if cuda else ["CPUExecutionProvider"]
515
+ session = onnxruntime.InferenceSession(w, providers=providers)
516
+ output_names = [x.name for x in session.get_outputs()]
517
+ meta = session.get_modelmeta().custom_metadata_map # metadata
518
+ if "stride" in meta:
519
+ stride, names = int(meta["stride"]), eval(meta["names"])
520
+ elif xml: # OpenVINO
521
+ LOGGER.info(f"Loading {w} for OpenVINO inference...")
522
+ check_requirements("openvino>=2023.0") # requires openvino-dev: https://pypi.org/project/openvino-dev/
523
+ from openvino.runtime import Core, Layout, get_batch
524
+
525
+ core = Core()
526
+ if not Path(w).is_file(): # if not *.xml
527
+ w = next(Path(w).glob("*.xml")) # get *.xml file from *_openvino_model dir
528
+ ov_model = core.read_model(model=w, weights=Path(w).with_suffix(".bin"))
529
+ if ov_model.get_parameters()[0].get_layout().empty:
530
+ ov_model.get_parameters()[0].set_layout(Layout("NCHW"))
531
+ batch_dim = get_batch(ov_model)
532
+ if batch_dim.is_static:
533
+ batch_size = batch_dim.get_length()
534
+ ov_compiled_model = core.compile_model(ov_model, device_name="AUTO") # AUTO selects best available device
535
+ stride, names = self._load_metadata(Path(w).with_suffix(".yaml")) # load metadata
536
+ elif engine: # TensorRT
537
+ LOGGER.info(f"Loading {w} for TensorRT inference...")
538
+ import tensorrt as trt # https://developer.nvidia.com/nvidia-tensorrt-download
539
+
540
+ check_version(trt.__version__, "7.0.0", hard=True) # require tensorrt>=7.0.0
541
+ if device.type == "cpu":
542
+ device = torch.device("cuda:0")
543
+ Binding = namedtuple("Binding", ("name", "dtype", "shape", "data", "ptr"))
544
+ logger = trt.Logger(trt.Logger.INFO)
545
+ with open(w, "rb") as f, trt.Runtime(logger) as runtime:
546
+ model = runtime.deserialize_cuda_engine(f.read())
547
+ context = model.create_execution_context()
548
+ bindings = OrderedDict()
549
+ output_names = []
550
+ fp16 = False # default updated below
551
+ dynamic = False
552
+ is_trt10 = not hasattr(model, "num_bindings")
553
+ num = range(model.num_io_tensors) if is_trt10 else range(model.num_bindings)
554
+ for i in num:
555
+ if is_trt10:
556
+ name = model.get_tensor_name(i)
557
+ dtype = trt.nptype(model.get_tensor_dtype(name))
558
+ is_input = model.get_tensor_mode(name) == trt.TensorIOMode.INPUT
559
+ if is_input:
560
+ if -1 in tuple(model.get_tensor_shape(name)): # dynamic
561
+ dynamic = True
562
+ context.set_input_shape(name, tuple(model.get_profile_shape(name, 0)[2]))
563
+ if dtype == np.float16:
564
+ fp16 = True
565
+ else: # output
566
+ output_names.append(name)
567
+ shape = tuple(context.get_tensor_shape(name))
568
+ else:
569
+ name = model.get_binding_name(i)
570
+ dtype = trt.nptype(model.get_binding_dtype(i))
571
+ if model.binding_is_input(i):
572
+ if -1 in tuple(model.get_binding_shape(i)): # dynamic
573
+ dynamic = True
574
+ context.set_binding_shape(i, tuple(model.get_profile_shape(0, i)[2]))
575
+ if dtype == np.float16:
576
+ fp16 = True
577
+ else: # output
578
+ output_names.append(name)
579
+ shape = tuple(context.get_binding_shape(i))
580
+ im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)
581
+ bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))
582
+ binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
583
+ batch_size = bindings["images"].shape[0] # if dynamic, this is instead max batch size
584
+ elif coreml: # CoreML
585
+ LOGGER.info(f"Loading {w} for CoreML inference...")
586
+ import coremltools as ct
587
+
588
+ model = ct.models.MLModel(w)
589
+ elif saved_model: # TF SavedModel
590
+ LOGGER.info(f"Loading {w} for TensorFlow SavedModel inference...")
591
+ import tensorflow as tf
592
+
593
+ keras = False # assume TF1 saved_model
594
+ model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)
595
+ elif pb: # GraphDef https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
596
+ LOGGER.info(f"Loading {w} for TensorFlow GraphDef inference...")
597
+ import tensorflow as tf
598
+
599
+ def wrap_frozen_graph(gd, inputs, outputs):
600
+ """Wraps a TensorFlow GraphDef for inference, returning a pruned function."""
601
+ x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) # wrapped
602
+ ge = x.graph.as_graph_element
603
+ return x.prune(tf.nest.map_structure(ge, inputs), tf.nest.map_structure(ge, outputs))
604
+
605
+ def gd_outputs(gd):
606
+ """Generates a sorted list of graph outputs excluding NoOp nodes and inputs, formatted as '<name>:0'."""
607
+ name_list, input_list = [], []
608
+ for node in gd.node: # tensorflow.core.framework.node_def_pb2.NodeDef
609
+ name_list.append(node.name)
610
+ input_list.extend(node.input)
611
+ return sorted(f"{x}:0" for x in list(set(name_list) - set(input_list)) if not x.startswith("NoOp"))
612
+
613
+ gd = tf.Graph().as_graph_def() # TF GraphDef
614
+ with open(w, "rb") as f:
615
+ gd.ParseFromString(f.read())
616
+ frozen_func = wrap_frozen_graph(gd, inputs="x:0", outputs=gd_outputs(gd))
617
+ elif tflite or edgetpu: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
618
+ try: # https://coral.ai/docs/edgetpu/tflite-python/#update-existing-tf-lite-code-for-the-edge-tpu
619
+ from tflite_runtime.interpreter import Interpreter, load_delegate
620
+ except ImportError:
621
+ import tensorflow as tf
622
+
623
+ Interpreter, load_delegate = (
624
+ tf.lite.Interpreter,
625
+ tf.lite.experimental.load_delegate,
626
+ )
627
+ if edgetpu: # TF Edge TPU https://coral.ai/software/#edgetpu-runtime
628
+ LOGGER.info(f"Loading {w} for TensorFlow Lite Edge TPU inference...")
629
+ delegate = {"Linux": "libedgetpu.so.1", "Darwin": "libedgetpu.1.dylib", "Windows": "edgetpu.dll"}[
630
+ platform.system()
631
+ ]
632
+ interpreter = Interpreter(model_path=w, experimental_delegates=[load_delegate(delegate)])
633
+ else: # TFLite
634
+ LOGGER.info(f"Loading {w} for TensorFlow Lite inference...")
635
+ interpreter = Interpreter(model_path=w) # load TFLite model
636
+ interpreter.allocate_tensors() # allocate
637
+ input_details = interpreter.get_input_details() # inputs
638
+ output_details = interpreter.get_output_details() # outputs
639
+ # load metadata
640
+ with contextlib.suppress(zipfile.BadZipFile):
641
+ with zipfile.ZipFile(w, "r") as model:
642
+ meta_file = model.namelist()[0]
643
+ meta = ast.literal_eval(model.read(meta_file).decode("utf-8"))
644
+ stride, names = int(meta["stride"]), meta["names"]
645
+ elif tfjs: # TF.js
646
+ raise NotImplementedError("ERROR: YOLOv5 TF.js inference is not supported")
647
+ elif paddle: # PaddlePaddle
648
+ LOGGER.info(f"Loading {w} for PaddlePaddle inference...")
649
+ check_requirements("paddlepaddle-gpu" if cuda else "paddlepaddle")
650
+ import paddle.inference as pdi
651
+
652
+ if not Path(w).is_file(): # if not *.pdmodel
653
+ w = next(Path(w).rglob("*.pdmodel")) # get *.pdmodel file from *_paddle_model dir
654
+ weights = Path(w).with_suffix(".pdiparams")
655
+ config = pdi.Config(str(w), str(weights))
656
+ if cuda:
657
+ config.enable_use_gpu(memory_pool_init_size_mb=2048, device_id=0)
658
+ predictor = pdi.create_predictor(config)
659
+ input_handle = predictor.get_input_handle(predictor.get_input_names()[0])
660
+ output_names = predictor.get_output_names()
661
+ elif triton: # NVIDIA Triton Inference Server
662
+ LOGGER.info(f"Using {w} as Triton Inference Server...")
663
+ check_requirements("tritonclient[all]")
664
+ from utils.triton import TritonRemoteModel
665
+
666
+ model = TritonRemoteModel(url=w)
667
+ nhwc = model.runtime.startswith("tensorflow")
668
+ else:
669
+ raise NotImplementedError(f"ERROR: {w} is not a supported format")
670
+
671
+ # class names
672
+ if "names" not in locals():
673
+ names = yaml_load(data)["names"] if data else {i: f"class{i}" for i in range(999)}
674
+ if names[0] == "n01440764" and len(names) == 1000: # ImageNet
675
+ names = yaml_load(ROOT / "data/ImageNet.yaml")["names"] # human-readable names
676
+
677
+ self.__dict__.update(locals()) # assign all variables to self
678
+
679
+ def forward(self, im, augment=False, visualize=False):
680
+ """Performs YOLOv5 inference on input images with options for augmentation and visualization."""
681
+ b, ch, h, w = im.shape # batch, channel, height, width
682
+ if self.fp16 and im.dtype != torch.float16:
683
+ im = im.half() # to FP16
684
+ if self.nhwc:
685
+ im = im.permute(0, 2, 3, 1) # torch BCHW to numpy BHWC shape(1,320,192,3)
686
+
687
+ if self.pt: # PyTorch
688
+ y = self.model(im, augment=augment, visualize=visualize) if augment or visualize else self.model(im)
689
+ elif self.jit: # TorchScript
690
+ y = self.model(im)
691
+ elif self.dnn: # ONNX OpenCV DNN
692
+ im = im.cpu().numpy() # torch to numpy
693
+ self.net.setInput(im)
694
+ y = self.net.forward()
695
+ elif self.onnx: # ONNX Runtime
696
+ im = im.cpu().numpy() # torch to numpy
697
+ y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})
698
+ elif self.xml: # OpenVINO
699
+ im = im.cpu().numpy() # FP32
700
+ y = list(self.ov_compiled_model(im).values())
701
+ elif self.engine: # TensorRT
702
+ if self.dynamic and im.shape != self.bindings["images"].shape:
703
+ i = self.model.get_binding_index("images")
704
+ self.context.set_binding_shape(i, im.shape) # reshape if dynamic
705
+ self.bindings["images"] = self.bindings["images"]._replace(shape=im.shape)
706
+ for name in self.output_names:
707
+ i = self.model.get_binding_index(name)
708
+ self.bindings[name].data.resize_(tuple(self.context.get_binding_shape(i)))
709
+ s = self.bindings["images"].shape
710
+ assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
711
+ self.binding_addrs["images"] = int(im.data_ptr())
712
+ self.context.execute_v2(list(self.binding_addrs.values()))
713
+ y = [self.bindings[x].data for x in sorted(self.output_names)]
714
+ elif self.coreml: # CoreML
715
+ im = im.cpu().numpy()
716
+ im = Image.fromarray((im[0] * 255).astype("uint8"))
717
+ # im = im.resize((192, 320), Image.BILINEAR)
718
+ y = self.model.predict({"image": im}) # coordinates are xywh normalized
719
+ if "confidence" in y:
720
+ box = xywh2xyxy(y["coordinates"] * [[w, h, w, h]]) # xyxy pixels
721
+ conf, cls = y["confidence"].max(1), y["confidence"].argmax(1).astype(np.float)
722
+ y = np.concatenate((box, conf.reshape(-1, 1), cls.reshape(-1, 1)), 1)
723
+ else:
724
+ y = list(reversed(y.values())) # reversed for segmentation models (pred, proto)
725
+ elif self.paddle: # PaddlePaddle
726
+ im = im.cpu().numpy().astype(np.float32)
727
+ self.input_handle.copy_from_cpu(im)
728
+ self.predictor.run()
729
+ y = [self.predictor.get_output_handle(x).copy_to_cpu() for x in self.output_names]
730
+ elif self.triton: # NVIDIA Triton Inference Server
731
+ y = self.model(im)
732
+ else: # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
733
+ im = im.cpu().numpy()
734
+ if self.saved_model: # SavedModel
735
+ y = self.model(im, training=False) if self.keras else self.model(im)
736
+ elif self.pb: # GraphDef
737
+ y = self.frozen_func(x=self.tf.constant(im))
738
+ else: # Lite or Edge TPU
739
+ input = self.input_details[0]
740
+ int8 = input["dtype"] == np.uint8 # is TFLite quantized uint8 model
741
+ if int8:
742
+ scale, zero_point = input["quantization"]
743
+ im = (im / scale + zero_point).astype(np.uint8) # de-scale
744
+ self.interpreter.set_tensor(input["index"], im)
745
+ self.interpreter.invoke()
746
+ y = []
747
+ for output in self.output_details:
748
+ x = self.interpreter.get_tensor(output["index"])
749
+ if int8:
750
+ scale, zero_point = output["quantization"]
751
+ x = (x.astype(np.float32) - zero_point) * scale # re-scale
752
+ y.append(x)
753
+ y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
754
+ y[0][..., :4] *= [w, h, w, h] # xywh normalized to pixels
755
+
756
+ if isinstance(y, (list, tuple)):
757
+ return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y]
758
+ else:
759
+ return self.from_numpy(y)
760
+
761
+ def from_numpy(self, x):
762
+ """Converts a NumPy array to a torch tensor, maintaining device compatibility."""
763
+ return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x
764
+
765
+ def warmup(self, imgsz=(1, 3, 640, 640)):
766
+ """Performs a single inference warmup to initialize model weights, accepting an `imgsz` tuple for image size."""
767
+ warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton
768
+ if any(warmup_types) and (self.device.type != "cpu" or self.triton):
769
+ im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input
770
+ for _ in range(2 if self.jit else 1): #
771
+ self.forward(im) # warmup
772
+
773
+ @staticmethod
774
+ def _model_type(p="path/to/model.pt"):
775
+ """
776
+ Determines model type from file path or URL, supporting various export formats.
777
+
778
+ Example: path='path/to/model.onnx' -> type=onnx
779
+ """
780
+ # types = [pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle]
781
+ from export import export_formats
782
+ from utils.downloads import is_url
783
+
784
+ sf = list(export_formats().Suffix) # export suffixes
785
+ if not is_url(p, check=False):
786
+ check_suffix(p, sf) # checks
787
+ url = urlparse(p) # if url may be Triton inference server
788
+ types = [s in Path(p).name for s in sf]
789
+ types[8] &= not types[9] # tflite &= not edgetpu
790
+ triton = not any(types) and all([any(s in url.scheme for s in ["http", "grpc"]), url.netloc])
791
+ return types + [triton]
792
+
793
+ @staticmethod
794
+ def _load_metadata(f=Path("path/to/meta.yaml")):
795
+ """Loads metadata from a YAML file, returning strides and names if the file exists, otherwise `None`."""
796
+ if f.exists():
797
+ d = yaml_load(f)
798
+ return d["stride"], d["names"] # assign stride, names
799
+ return None, None
800
+
801
+
802
+ class AutoShape(nn.Module):
803
+ """AutoShape class for robust YOLOv5 inference with preprocessing, NMS, and support for various input formats."""
804
+
805
+ conf = 0.25 # NMS confidence threshold
806
+ iou = 0.45 # NMS IoU threshold
807
+ agnostic = False # NMS class-agnostic
808
+ multi_label = False # NMS multiple labels per box
809
+ classes = None # (optional list) filter by class, i.e. = [0, 15, 16] for COCO persons, cats and dogs
810
+ max_det = 1000 # maximum number of detections per image
811
+ amp = False # Automatic Mixed Precision (AMP) inference
812
+
813
+ def __init__(self, model, verbose=True):
814
+ """Initializes YOLOv5 model for inference, setting up attributes and preparing model for evaluation."""
815
+ super().__init__()
816
+ if verbose:
817
+ LOGGER.info("Adding AutoShape... ")
818
+ copy_attr(self, model, include=("yaml", "nc", "hyp", "names", "stride", "abc"), exclude=()) # copy attributes
819
+ self.dmb = isinstance(model, DetectMultiBackend) # DetectMultiBackend() instance
820
+ self.pt = not self.dmb or model.pt # PyTorch model
821
+ self.model = model.eval()
822
+ if self.pt:
823
+ m = self.model.model.model[-1] if self.dmb else self.model.model[-1] # Detect()
824
+ m.inplace = False # Detect.inplace=False for safe multithread inference
825
+ m.export = True # do not output loss values
826
+
827
+ def _apply(self, fn):
828
+ """
829
+ Applies to(), cpu(), cuda(), half() etc.
830
+
831
+ to model tensors excluding parameters or registered buffers.
832
+ """
833
+ self = super()._apply(fn)
834
+ if self.pt:
835
+ m = self.model.model.model[-1] if self.dmb else self.model.model[-1] # Detect()
836
+ m.stride = fn(m.stride)
837
+ m.grid = list(map(fn, m.grid))
838
+ if isinstance(m.anchor_grid, list):
839
+ m.anchor_grid = list(map(fn, m.anchor_grid))
840
+ return self
841
+
842
+ @smart_inference_mode()
843
+ def forward(self, ims, size=640, augment=False, profile=False):
844
+ """
845
+ Performs inference on inputs with optional augment & profiling.
846
+
847
+ Supports various formats including file, URI, OpenCV, PIL, numpy, torch.
848
+ """
849
+ # For size(height=640, width=1280), RGB images example inputs are:
850
+ # file: ims = 'data/images/zidane.jpg' # str or PosixPath
851
+ # URI: = 'https://ultralytics.com/images/zidane.jpg'
852
+ # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
853
+ # PIL: = Image.open('image.jpg') or ImageGrab.grab() # HWC x(640,1280,3)
854
+ # numpy: = np.zeros((640,1280,3)) # HWC
855
+ # torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)
856
+ # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
857
+
858
+ dt = (Profile(), Profile(), Profile())
859
+ with dt[0]:
860
+ if isinstance(size, int): # expand
861
+ size = (size, size)
862
+ p = next(self.model.parameters()) if self.pt else torch.empty(1, device=self.model.device) # param
863
+ autocast = self.amp and (p.device.type != "cpu") # Automatic Mixed Precision (AMP) inference
864
+ if isinstance(ims, torch.Tensor): # torch
865
+ with amp.autocast(autocast):
866
+ return self.model(ims.to(p.device).type_as(p), augment=augment) # inference
867
+
868
+ # Pre-process
869
+ n, ims = (len(ims), list(ims)) if isinstance(ims, (list, tuple)) else (1, [ims]) # number, list of images
870
+ shape0, shape1, files = [], [], [] # image and inference shapes, filenames
871
+ for i, im in enumerate(ims):
872
+ f = f"image{i}" # filename
873
+ if isinstance(im, (str, Path)): # filename or uri
874
+ im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith("http") else im), im
875
+ im = np.asarray(exif_transpose(im))
876
+ elif isinstance(im, Image.Image): # PIL Image
877
+ im, f = np.asarray(exif_transpose(im)), getattr(im, "filename", f) or f
878
+ files.append(Path(f).with_suffix(".jpg").name)
879
+ if im.shape[0] < 5: # image in CHW
880
+ im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
881
+ im = im[..., :3] if im.ndim == 3 else cv2.cvtColor(im, cv2.COLOR_GRAY2BGR) # enforce 3ch input
882
+ s = im.shape[:2] # HWC
883
+ shape0.append(s) # image shape
884
+ g = max(size) / max(s) # gain
885
+ shape1.append([int(y * g) for y in s])
886
+ ims[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
887
+ shape1 = [make_divisible(x, self.stride) for x in np.array(shape1).max(0)] # inf shape
888
+ x = [letterbox(im, shape1, auto=False)[0] for im in ims] # pad
889
+ x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2))) # stack and BHWC to BCHW
890
+ x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
891
+
892
+ with amp.autocast(autocast):
893
+ # Inference
894
+ with dt[1]:
895
+ y = self.model(x, augment=augment) # forward
896
+
897
+ # Post-process
898
+ with dt[2]:
899
+ y = non_max_suppression(
900
+ y if self.dmb else y[0],
901
+ self.conf,
902
+ self.iou,
903
+ self.classes,
904
+ self.agnostic,
905
+ self.multi_label,
906
+ max_det=self.max_det,
907
+ ) # NMS
908
+ for i in range(n):
909
+ scale_boxes(shape1, y[i][:, :4], shape0[i])
910
+
911
+ return Detections(ims, y, files, dt, self.names, x.shape)
912
+
913
+
914
+ class Detections:
915
+ """Manages YOLOv5 detection results with methods for visualization, saving, cropping, and exporting detections."""
916
+
917
+ def __init__(self, ims, pred, files, times=(0, 0, 0), names=None, shape=None):
918
+ """Initializes the YOLOv5 Detections class with image info, predictions, filenames, timing and normalization."""
919
+ super().__init__()
920
+ d = pred[0].device # device
921
+ gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in ims] # normalizations
922
+ self.ims = ims # list of images as numpy arrays
923
+ self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
924
+ self.names = names # class names
925
+ self.files = files # image filenames
926
+ self.times = times # profiling times
927
+ self.xyxy = pred # xyxy pixels
928
+ self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
929
+ self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
930
+ self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
931
+ self.n = len(self.pred) # number of images (batch size)
932
+ self.t = tuple(x.t / self.n * 1e3 for x in times) # timestamps (ms)
933
+ self.s = tuple(shape) # inference BCHW shape
934
+
935
+ def _run(self, pprint=False, show=False, save=False, crop=False, render=False, labels=True, save_dir=Path("")):
936
+ """Executes model predictions, displaying and/or saving outputs with optional crops and labels."""
937
+ s, crops = "", []
938
+ for i, (im, pred) in enumerate(zip(self.ims, self.pred)):
939
+ s += f"\nimage {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} " # string
940
+ if pred.shape[0]:
941
+ for c in pred[:, -1].unique():
942
+ n = (pred[:, -1] == c).sum() # detections per class
943
+ s += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string
944
+ s = s.rstrip(", ")
945
+ if show or save or render or crop:
946
+ annotator = Annotator(im, example=str(self.names))
947
+ for *box, conf, cls in reversed(pred): # xyxy, confidence, class
948
+ label = f"{self.names[int(cls)]} {conf:.2f}"
949
+ if crop:
950
+ file = save_dir / "crops" / self.names[int(cls)] / self.files[i] if save else None
951
+ crops.append(
952
+ {
953
+ "box": box,
954
+ "conf": conf,
955
+ "cls": cls,
956
+ "label": label,
957
+ "im": save_one_box(box, im, file=file, save=save),
958
+ }
959
+ )
960
+ else: # all others
961
+ annotator.box_label(box, label if labels else "", color=colors(cls))
962
+ im = annotator.im
963
+ else:
964
+ s += "(no detections)"
965
+
966
+ im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np
967
+ if show:
968
+ if is_jupyter():
969
+ from IPython.display import display
970
+
971
+ display(im)
972
+ else:
973
+ im.show(self.files[i])
974
+ if save:
975
+ f = self.files[i]
976
+ im.save(save_dir / f) # save
977
+ if i == self.n - 1:
978
+ LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
979
+ if render:
980
+ self.ims[i] = np.asarray(im)
981
+ if pprint:
982
+ s = s.lstrip("\n")
983
+ return f"{s}\nSpeed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {self.s}" % self.t
984
+ if crop:
985
+ if save:
986
+ LOGGER.info(f"Saved results to {save_dir}\n")
987
+ return crops
988
+
989
+ @TryExcept("Showing images is not supported in this environment")
990
+ def show(self, labels=True):
991
+ """
992
+ Displays detection results with optional labels.
993
+
994
+ Usage: show(labels=True)
995
+ """
996
+ self._run(show=True, labels=labels) # show results
997
+
998
+ def save(self, labels=True, save_dir="runs/detect/exp", exist_ok=False):
999
+ """
1000
+ Saves detection results with optional labels to a specified directory.
1001
+
1002
+ Usage: save(labels=True, save_dir='runs/detect/exp', exist_ok=False)
1003
+ """
1004
+ save_dir = increment_path(save_dir, exist_ok, mkdir=True) # increment save_dir
1005
+ self._run(save=True, labels=labels, save_dir=save_dir) # save results
1006
+
1007
+ def crop(self, save=True, save_dir="runs/detect/exp", exist_ok=False):
1008
+ """
1009
+ Crops detection results, optionally saves them to a directory.
1010
+
1011
+ Args: save (bool), save_dir (str), exist_ok (bool).
1012
+ """
1013
+ save_dir = increment_path(save_dir, exist_ok, mkdir=True) if save else None
1014
+ return self._run(crop=True, save=save, save_dir=save_dir) # crop results
1015
+
1016
+ def render(self, labels=True):
1017
+ """Renders detection results with optional labels on images; args: labels (bool) indicating label inclusion."""
1018
+ self._run(render=True, labels=labels) # render results
1019
+ return self.ims
1020
+
1021
+ def pandas(self):
1022
+ """
1023
+ Returns detections as pandas DataFrames for various box formats (xyxy, xyxyn, xywh, xywhn).
1024
+
1025
+ Example: print(results.pandas().xyxy[0]).
1026
+ """
1027
+ new = copy(self) # return copy
1028
+ ca = "xmin", "ymin", "xmax", "ymax", "confidence", "class", "name" # xyxy columns
1029
+ cb = "xcenter", "ycenter", "width", "height", "confidence", "class", "name" # xywh columns
1030
+ for k, c in zip(["xyxy", "xyxyn", "xywh", "xywhn"], [ca, ca, cb, cb]):
1031
+ a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update
1032
+ setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
1033
+ return new
1034
+
1035
+ def tolist(self):
1036
+ """
1037
+ Converts a Detections object into a list of individual detection results for iteration.
1038
+
1039
+ Example: for result in results.tolist():
1040
+ """
1041
+ r = range(self.n) # iterable
1042
+ return [
1043
+ Detections(
1044
+ [self.ims[i]],
1045
+ [self.pred[i]],
1046
+ [self.files[i]],
1047
+ self.times,
1048
+ self.names,
1049
+ self.s,
1050
+ )
1051
+ for i in r
1052
+ ]
1053
+
1054
+ def print(self):
1055
+ """Logs the string representation of the current object's state via the LOGGER."""
1056
+ LOGGER.info(self.__str__())
1057
+
1058
+ def __len__(self):
1059
+ """Returns the number of results stored, overrides the default len(results)."""
1060
+ return self.n
1061
+
1062
+ def __str__(self):
1063
+ """Returns a string representation of the model's results, suitable for printing, overrides default
1064
+ print(results).
1065
+ """
1066
+ return self._run(pprint=True) # print results
1067
+
1068
+ def __repr__(self):
1069
+ """Returns a string representation of the YOLOv5 object, including its class and formatted results."""
1070
+ return f"YOLOv5 {self.__class__} instance\n" + self.__str__()
1071
+
1072
+
1073
+ class Proto(nn.Module):
1074
+ """YOLOv5 mask Proto module for segmentation models, performing convolutions and upsampling on input tensors."""
1075
+
1076
+ def __init__(self, c1, c_=256, c2=32):
1077
+ """Initializes YOLOv5 Proto module for segmentation with input, proto, and mask channels configuration."""
1078
+ super().__init__()
1079
+ self.cv1 = Conv(c1, c_, k=3)
1080
+ self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
1081
+ self.cv2 = Conv(c_, c_, k=3)
1082
+ self.cv3 = Conv(c_, c2)
1083
+
1084
+ def forward(self, x):
1085
+ """Performs a forward pass using convolutional layers and upsampling on input tensor `x`."""
1086
+ return self.cv3(self.cv2(self.upsample(self.cv1(x))))
1087
+
1088
+
1089
+ class Classify(nn.Module):
1090
+ """YOLOv5 classification head with convolution, pooling, and dropout layers for channel transformation."""
1091
+
1092
+ def __init__(
1093
+ self, c1, c2, k=1, s=1, p=None, g=1, dropout_p=0.0
1094
+ ): # ch_in, ch_out, kernel, stride, padding, groups, dropout probability
1095
+ """Initializes YOLOv5 classification head with convolution, pooling, and dropout layers for input to output
1096
+ channel transformation.
1097
+ """
1098
+ super().__init__()
1099
+ c_ = 1280 # efficientnet_b0 size
1100
+ self.conv = Conv(c1, c_, k, s, autopad(k, p), g)
1101
+ self.pool = nn.AdaptiveAvgPool2d(1) # to x(b,c_,1,1)
1102
+ self.drop = nn.Dropout(p=dropout_p, inplace=True)
1103
+ self.linear = nn.Linear(c_, c2) # to x(b,c2)
1104
+
1105
+ def forward(self, x):
1106
+ """Processes input through conv, pool, drop, and linear layers; supports list concatenation input."""
1107
+ if isinstance(x, list):
1108
+ x = torch.cat(x, 1)
1109
+ return self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))