floratest1
/
SpireCV
1
0
Fork 1
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix model sync

This commit is contained in:
jario 2024-01-19 23:56:05 +08:00
parent 9c4d8a400c
commit 83bf2f9656
1 changed files with 11 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
scripts/model_sync.py
View File

@ -52,14 +52,18 @@ def main():
print("[1] Model:", line, "EXIST!") print("[1] Model:", line, "EXIST!")
if line.startswith('Nv'): if line.startswith('Nv'):
net = line.split('-')[2] net = line.split('-')[2]
if len(net.split('_')) == 3: if net.startswith("yolov5"):
name, seg, ncls = net.split('_') if len(net.split('_')) == 3:
engine_fn = os.path.splitext(model_file)[0].replace(net, name + "_" + seg + '_b1_' + ncls) + '.engine' name, seg, ncls = net.split('_')
online_fn = os.path.splitext(model_file)[0].replace(net, name + "_" + seg + '_b1_' + ncls) + '-online.engine' engine_fn = os.path.splitext(model_file)[0].replace(net, name + "_" + seg + '_b1_' + ncls) + '.engine'
online_fn = os.path.splitext(model_file)[0].replace(net, name + "_" + seg + '_b1_' + ncls) + '-online.engine'
else:
name, ncls = net.split('_')
engine_fn = os.path.splitext(model_file)[0].replace(net, name + '_b1_' + ncls) + '.engine'
online_fn = os.path.splitext(model_file)[0].replace(net, name + '_b1_' + ncls) + '-online.engine'
else: else:
name, ncls = net.split('_') engine_fn = os.path.splitext(model_file)[0] + '.engine'
engine_fn = os.path.splitext(model_file)[0].replace(net, name + '_b1_' + ncls) + '.engine' online_fn = os.path.splitext(model_file)[0] + '-online.engine'
online_fn = os.path.splitext(model_file)[0].replace(net, name + '_b1_' + ncls) + '-online.engine'
if not os.path.exists(engine_fn) and not os.path.exists(online_fn): if not os.path.exists(engine_fn) and not os.path.exists(online_fn):
if net.startswith("yolov5"): if net.startswith("yolov5"):
if len(net.split('_')) == 3: if len(net.split('_')) == 3: