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

Fix gen_wts_yoloV5.wts

Browse Source
This commit is contained in:
Marcos Luciano
2022-07-21 11:30:20 -03:00
parent 9b753352c8
commit d09879d557
1 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
utils/gen_wts_yoloV5.py
View File

@@ -173,9 +173,9 @@ class Layers(object):
self.fc.write('\n[convolutional]\n' + self.fc.write('\n[convolutional]\n' +
b + b +
'filters=%d\n' % filters + 'filters=%d\n' % filters +
'size=%s\n' % (size[0] if type(size) != int and size[0] == size[1] else str(size)[1:-1]) + 'size=%s\n' % self.get_value(size) +
'stride=%s\n' % (stride[0] if type(stride) != int and stride[0] == stride[1] else str(stride)[1:-1]) + 'stride=%s\n' % self.get_value(stride) +
'pad=%s\n' % (pad[0] if type(pad) != int and pad[0] == pad[1] else str(pad)[1:-1]) + 'pad=%s\n' % self.get_value(pad) +
g + g +
w + w +
'activation=%s\n' % act) 'activation=%s\n' % act)
@@ -265,6 +265,11 @@ class Layers(object):
n += 1 n += 1
self.masks.append(str(mask)[1:-1]) self.masks.append(str(mask)[1:-1])
def get_value(self, key):
if type(key) == int:
return key
return key[0] if key[0] == key[1] else str(key)[1:-1]
def get_route(self, n): def get_route(self, n):
r = 0 r = 0
for i, b in enumerate(self.blocks): for i, b in enumerate(self.blocks):