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-03 12:55:46 -03:00
parent cd3ab90867
commit 842a22f6b9
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
utils/gen_wts_yoloV8.py
View File

@@ -3,6 +3,7 @@ import os
import struct import struct
import torch import torch
from ultralytics.yolo.utils.torch_utils import select_device from ultralytics.yolo.utils.torch_utils import select_device
from ultralytics.yolo.utils.tal import make_anchors
class Layers(object): class Layers(object):
@@ -306,7 +307,15 @@ with open(wts_file, 'w') as fw, open(cfg_file, 'w') as fc:
layers.Concat(child) layers.Concat(child)
elif child._get_name() == 'Detect': elif child._get_name() == 'Detect':
layers.Detect(child) layers.Detect(child)
layers.get_anchors(child.anchors.reshape([-1]), child.strides.reshape([-1])) 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))
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: else:
raise SystemExit('Model not supported') raise SystemExit('Model not supported')