bplusplus 1.1.0__py3-none-any.whl → 1.2.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.
- bplusplus/__init__.py +4 -2
- bplusplus/collect.py +69 -5
- bplusplus/hierarchical/test.py +670 -0
- bplusplus/hierarchical/train.py +676 -0
- bplusplus/prepare.py +228 -64
- bplusplus/resnet/test.py +473 -0
- bplusplus/resnet/train.py +329 -0
- bplusplus-1.2.0.dist-info/METADATA +249 -0
- bplusplus-1.2.0.dist-info/RECORD +12 -0
- bplusplus/yolov5detect/__init__.py +0 -1
- bplusplus/yolov5detect/detect.py +0 -444
- bplusplus/yolov5detect/export.py +0 -1530
- bplusplus/yolov5detect/insect.yaml +0 -8
- bplusplus/yolov5detect/models/__init__.py +0 -0
- bplusplus/yolov5detect/models/common.py +0 -1109
- bplusplus/yolov5detect/models/experimental.py +0 -130
- bplusplus/yolov5detect/models/hub/anchors.yaml +0 -56
- bplusplus/yolov5detect/models/hub/yolov3-spp.yaml +0 -52
- bplusplus/yolov5detect/models/hub/yolov3-tiny.yaml +0 -42
- bplusplus/yolov5detect/models/hub/yolov3.yaml +0 -52
- bplusplus/yolov5detect/models/hub/yolov5-bifpn.yaml +0 -49
- bplusplus/yolov5detect/models/hub/yolov5-fpn.yaml +0 -43
- bplusplus/yolov5detect/models/hub/yolov5-p2.yaml +0 -55
- bplusplus/yolov5detect/models/hub/yolov5-p34.yaml +0 -42
- bplusplus/yolov5detect/models/hub/yolov5-p6.yaml +0 -57
- bplusplus/yolov5detect/models/hub/yolov5-p7.yaml +0 -68
- bplusplus/yolov5detect/models/hub/yolov5-panet.yaml +0 -49
- bplusplus/yolov5detect/models/hub/yolov5l6.yaml +0 -61
- bplusplus/yolov5detect/models/hub/yolov5m6.yaml +0 -61
- bplusplus/yolov5detect/models/hub/yolov5n6.yaml +0 -61
- bplusplus/yolov5detect/models/hub/yolov5s-LeakyReLU.yaml +0 -50
- bplusplus/yolov5detect/models/hub/yolov5s-ghost.yaml +0 -49
- bplusplus/yolov5detect/models/hub/yolov5s-transformer.yaml +0 -49
- bplusplus/yolov5detect/models/hub/yolov5s6.yaml +0 -61
- bplusplus/yolov5detect/models/hub/yolov5x6.yaml +0 -61
- bplusplus/yolov5detect/models/segment/yolov5l-seg.yaml +0 -49
- bplusplus/yolov5detect/models/segment/yolov5m-seg.yaml +0 -49
- bplusplus/yolov5detect/models/segment/yolov5n-seg.yaml +0 -49
- bplusplus/yolov5detect/models/segment/yolov5s-seg.yaml +0 -49
- bplusplus/yolov5detect/models/segment/yolov5x-seg.yaml +0 -49
- bplusplus/yolov5detect/models/tf.py +0 -797
- bplusplus/yolov5detect/models/yolo.py +0 -495
- bplusplus/yolov5detect/models/yolov5l.yaml +0 -49
- bplusplus/yolov5detect/models/yolov5m.yaml +0 -49
- bplusplus/yolov5detect/models/yolov5n.yaml +0 -49
- bplusplus/yolov5detect/models/yolov5s.yaml +0 -49
- bplusplus/yolov5detect/models/yolov5x.yaml +0 -49
- bplusplus/yolov5detect/utils/__init__.py +0 -97
- bplusplus/yolov5detect/utils/activations.py +0 -134
- bplusplus/yolov5detect/utils/augmentations.py +0 -448
- bplusplus/yolov5detect/utils/autoanchor.py +0 -175
- bplusplus/yolov5detect/utils/autobatch.py +0 -70
- bplusplus/yolov5detect/utils/aws/__init__.py +0 -0
- bplusplus/yolov5detect/utils/aws/mime.sh +0 -26
- bplusplus/yolov5detect/utils/aws/resume.py +0 -41
- bplusplus/yolov5detect/utils/aws/userdata.sh +0 -27
- bplusplus/yolov5detect/utils/callbacks.py +0 -72
- bplusplus/yolov5detect/utils/dataloaders.py +0 -1385
- bplusplus/yolov5detect/utils/docker/Dockerfile +0 -73
- bplusplus/yolov5detect/utils/docker/Dockerfile-arm64 +0 -40
- bplusplus/yolov5detect/utils/docker/Dockerfile-cpu +0 -42
- bplusplus/yolov5detect/utils/downloads.py +0 -136
- bplusplus/yolov5detect/utils/flask_rest_api/README.md +0 -70
- bplusplus/yolov5detect/utils/flask_rest_api/example_request.py +0 -17
- bplusplus/yolov5detect/utils/flask_rest_api/restapi.py +0 -49
- bplusplus/yolov5detect/utils/general.py +0 -1294
- bplusplus/yolov5detect/utils/google_app_engine/Dockerfile +0 -25
- bplusplus/yolov5detect/utils/google_app_engine/additional_requirements.txt +0 -6
- bplusplus/yolov5detect/utils/google_app_engine/app.yaml +0 -16
- bplusplus/yolov5detect/utils/loggers/__init__.py +0 -476
- bplusplus/yolov5detect/utils/loggers/clearml/README.md +0 -222
- bplusplus/yolov5detect/utils/loggers/clearml/__init__.py +0 -0
- bplusplus/yolov5detect/utils/loggers/clearml/clearml_utils.py +0 -230
- bplusplus/yolov5detect/utils/loggers/clearml/hpo.py +0 -90
- bplusplus/yolov5detect/utils/loggers/comet/README.md +0 -250
- bplusplus/yolov5detect/utils/loggers/comet/__init__.py +0 -551
- bplusplus/yolov5detect/utils/loggers/comet/comet_utils.py +0 -151
- bplusplus/yolov5detect/utils/loggers/comet/hpo.py +0 -126
- bplusplus/yolov5detect/utils/loggers/comet/optimizer_config.json +0 -135
- bplusplus/yolov5detect/utils/loggers/wandb/__init__.py +0 -0
- bplusplus/yolov5detect/utils/loggers/wandb/wandb_utils.py +0 -210
- bplusplus/yolov5detect/utils/loss.py +0 -259
- bplusplus/yolov5detect/utils/metrics.py +0 -381
- bplusplus/yolov5detect/utils/plots.py +0 -517
- bplusplus/yolov5detect/utils/segment/__init__.py +0 -0
- bplusplus/yolov5detect/utils/segment/augmentations.py +0 -100
- bplusplus/yolov5detect/utils/segment/dataloaders.py +0 -366
- bplusplus/yolov5detect/utils/segment/general.py +0 -160
- bplusplus/yolov5detect/utils/segment/loss.py +0 -198
- bplusplus/yolov5detect/utils/segment/metrics.py +0 -225
- bplusplus/yolov5detect/utils/segment/plots.py +0 -152
- bplusplus/yolov5detect/utils/torch_utils.py +0 -482
- bplusplus/yolov5detect/utils/triton.py +0 -90
- bplusplus-1.1.0.dist-info/METADATA +0 -179
- bplusplus-1.1.0.dist-info/RECORD +0 -92
- {bplusplus-1.1.0.dist-info → bplusplus-1.2.0.dist-info}/LICENSE +0 -0
- {bplusplus-1.1.0.dist-info → bplusplus-1.2.0.dist-info}/WHEEL +0 -0
|
@@ -1,1109 +0,0 @@
|
|
|
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)))
|