thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix YOLOv8 conversion

Browse Source
This commit is contained in:
Marcos Luciano
2023-02-08 01:02:35 -03:00
parent 842a22f6b9
commit 940d244ad4
1 changed files with 5 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
utils/gen_wts_yoloV8.py
View File

@@ -307,13 +307,9 @@ with open(wts_file, 'w') as fw, open(cfg_file, 'w') as fc:
layers.Concat(child)
elif child._get_name() == 'Detect':
layers.Detect(child)
if child.anchors.nelement() > 0 and child.strides.nelement() > 0:
layers.get_anchors(child.anchors.reshape([-1]), child.strides.reshape([-1]))
else:
x = []
for stride in model.stride.tolist():
x.append(torch.zeros([1, 1, int(layers.height / stride), int(layers.width / stride)],
dtype=torch.float32))
x.append(torch.zeros([1, 1, int(layers.height / stride), int(layers.width / stride)], dtype=torch.float32))
anchor_points, stride_tensor = (x.transpose(0, 1) for x in make_anchors(x, child.stride, 0.5))
layers.get_anchors(anchor_points.reshape([-1]), stride_tensor.reshape([-1]))
else: