1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
| class Detect(nn.Module):
"""YOLO Detect head for object detection models.
This class implements the detection head used in YOLO models for predicting bounding boxes and class probabilities.
It supports both training and inference modes, with optional end-to-end detection capabilities.
Attributes:
dynamic (bool): Force grid reconstruction.
export (bool): Export mode flag.
format (str): Export format.
end2end (bool): End-to-end detection mode.
max_det (int): Maximum detections per image.
shape (tuple): Input shape.
anchors (torch.Tensor): Anchor points.
strides (torch.Tensor): Feature map strides.
legacy (bool): Backward compatibility for v3/v5/v8/v9 models.
xyxy (bool): Output format, xyxy or xywh.
nc (int): Number of classes.
nl (int): Number of detection layers.
reg_max (int): DFL channels.
no (int): Number of outputs per anchor.
stride (torch.Tensor): Strides computed during build.
cv2 (nn.ModuleList): Convolution layers for box regression.
cv3 (nn.ModuleList): Convolution layers for classification.
dfl (nn.Module): Distribution Focal Loss layer.
one2one_cv2 (nn.ModuleList): One-to-one convolution layers for box regression.
one2one_cv3 (nn.ModuleList): One-to-one convolution layers for classification.
Methods:
forward: Perform forward pass and return predictions.
forward_end2end: Perform forward pass for end-to-end detection.
bias_init: Initialize detection head biases.
decode_bboxes: Decode bounding boxes from predictions.
postprocess: Post-process model predictions.
Examples:
Create a detection head for 80 classes
>>> detect = Detect(nc=80, ch=(256, 512, 1024))
>>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
>>> outputs = detect(x)
"""
dynamic = False
export = False
format = None
max_det = 300
agnostic_nms = False
shape = None
anchors = torch.empty(0)
strides = torch.empty(0)
legacy = False
xyxy = False
def __init__(self, nc: int = 80, reg_max=16, end2end=False, ch: tuple = ()):
"""Initialize the YOLO detection layer with specified number of classes and channels.
Args:
nc (int): Number of classes.
reg_max (int): Maximum number of DFL channels.
end2end (bool): Whether to use end-to-end NMS-free detection.
ch (tuple): Tuple of channel sizes from backbone feature maps.
"""
super().__init__()
self.nc = nc
self.nl = len(ch)
self.reg_max = reg_max
self.no = nc + self.reg_max * 4
self.stride = torch.zeros(self.nl)
c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100))
self.cv2 = nn.ModuleList(
nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch
)
self.cv3 = (
nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch)
if self.legacy
else nn.ModuleList(
nn.Sequential(
nn.Sequential(DWConv(x, x, 3), Conv(x, c3, 1)),
nn.Sequential(DWConv(c3, c3, 3), Conv(c3, c3, 1)),
nn.Conv2d(c3, self.nc, 1),
)
for x in ch
)
)
self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()
if end2end:
self.one2one_cv2 = copy.deepcopy(self.cv2)
self.one2one_cv3 = copy.deepcopy(self.cv3)
@property
def one2many(self):
"""Returns the one-to-many head components, here for v5/v5/v8/v9/11 backward compatibility."""
return dict(box_head=self.cv2, cls_head=self.cv3)
@property
def one2one(self):
"""Returns the one-to-one head components."""
return dict(box_head=self.one2one_cv2, cls_head=self.one2one_cv3)
@property
def end2end(self):
"""Checks if the model has one2one for v5/v5/v8/v9/11 backward compatibility."""
return getattr(self, "_end2end", True) and hasattr(self, "one2one")
@end2end.setter
def end2end(self, value):
"""Override the end-to-end detection mode."""
self._end2end = value
def forward_head(
self, x: list[torch.Tensor], box_head: torch.nn.Module = None, cls_head: torch.nn.Module = None
) -> dict[str, torch.Tensor]:
"""Concatenates and returns predicted bounding boxes and class probabilities."""
if box_head is None or cls_head is None:
return dict()
bs = x[0].shape[0]
boxes = torch.cat([box_head[i](x[i]).view(bs, 4 * self.reg_max, -1) for i in range(self.nl)], dim=-1)
scores = torch.cat([cls_head[i](x[i]).view(bs, self.nc, -1) for i in range(self.nl)], dim=-1)
return dict(boxes=boxes, scores=scores, feats=x)
def forward(
self, x: list[torch.Tensor]
) -> dict[str, torch.Tensor] | torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
"""Concatenates and returns predicted bounding boxes and class probabilities."""
preds = self.forward_head(x, **self.one2many)
if self.end2end:
x_detach = [xi.detach() for xi in x]
one2one = self.forward_head(x_detach, **self.one2one)
preds = {"one2many": preds, "one2one": one2one}
if self.training:
return preds
y = self._inference(preds["one2one"] if self.end2end else preds)
if self.end2end:
y = self.postprocess(y.permute(0, 2, 1))
return y if self.export else (y, preds)
def _inference(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
"""Decode predicted bounding boxes and class probabilities based on multiple-level feature maps.
Args:
x (dict[str, torch.Tensor]): List of feature maps from different detection layers.
Returns:
(torch.Tensor): Concatenated tensor of decoded bounding boxes and class probabilities.
"""
dbox = self._get_decode_boxes(x)
return torch.cat((dbox, x["scores"].sigmoid()), 1)
def _get_decode_boxes(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
"""Get decoded boxes based on anchors and strides."""
shape = x["feats"][0].shape
if self.dynamic or self.shape != shape:
self.anchors, self.strides = (a.transpose(0, 1) for a in make_anchors(x["feats"], self.stride, 0.5))
self.shape = shape
dbox = self.decode_bboxes(self.dfl(x["boxes"]), self.anchors.unsqueeze(0)) * self.strides
return dbox
def bias_init(self):
"""Initialize Detect() biases, WARNING: requires stride availability."""
for i, (a, b) in enumerate(zip(self.one2many["box_head"], self.one2many["cls_head"])):
a[-1].bias.data[:] = 2.0
b[-1].bias.data[: self.nc] = math.log(
5 / self.nc / (640 / self.stride[i]) ** 2
)
if self.end2end:
for i, (a, b) in enumerate(zip(self.one2one["box_head"], self.one2one["cls_head"])):
a[-1].bias.data[:] = 2.0
b[-1].bias.data[: self.nc] = math.log(
5 / self.nc / (640 / self.stride[i]) ** 2
)
def decode_bboxes(self, bboxes: torch.Tensor, anchors: torch.Tensor, xywh: bool = True) -> torch.Tensor:
"""Decode bounding boxes from predictions."""
return dist2bbox(
bboxes,
anchors,
xywh=xywh and not self.end2end and not self.xyxy,
dim=1,
)
def postprocess(self, preds: torch.Tensor) -> torch.Tensor:
"""Post-processes YOLO model predictions.
Args:
preds (torch.Tensor): Raw predictions with shape (batch_size, num_anchors, 4 + nc) with last dimension
format [x, y, w, h, class_probs].
Returns:
(torch.Tensor): Processed predictions with shape (batch_size, min(max_det, num_anchors), 6) and last
dimension format [x, y, w, h, max_class_prob, class_index].
"""
boxes, scores = preds.split([4, self.nc], dim=-1)
scores, conf, idx = self.get_topk_index(scores, self.max_det)
boxes = boxes.gather(dim=1, index=idx.repeat(1, 1, 4))
return torch.cat([boxes, scores, conf], dim=-1)
def get_topk_index(self, scores: torch.Tensor, max_det: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Get top-k indices from scores.
Args:
scores (torch.Tensor): Scores tensor with shape (batch_size, num_anchors, num_classes).
max_det (int): Maximum detections per image.
Returns:
(torch.Tensor, torch.Tensor, torch.Tensor): Top scores, class indices, and filtered indices.
"""
batch_size, anchors, nc = scores.shape
k = max_det if self.export else min(max_det, anchors)
if self.agnostic_nms:
scores, labels = scores.max(dim=-1, keepdim=True)
scores, indices = scores.topk(k, dim=1)
labels = labels.gather(1, indices)
return scores, labels, indices
ori_index = scores.max(dim=-1)[0].topk(k)[1].unsqueeze(-1)
scores = scores.gather(dim=1, index=ori_index.repeat(1, 1, nc))
scores, index = scores.flatten(1).topk(k)
idx = ori_index[torch.arange(batch_size)[..., None], index // nc]
return scores[..., None], (index % nc)[..., None].float(), idx
def fuse(self) -> None:
"""Remove the one2many head for inference optimization."""
self.cv2 = self.cv3 = None
|