From 83bf2f9656272b61ec05e5e50ee381bb99d103f0 Mon Sep 17 00:00:00 2001 From: jario Date: Fri, 19 Jan 2024 23:56:05 +0800 Subject: [PATCH] fix model sync --- scripts/model_sync.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/scripts/model_sync.py b/scripts/model_sync.py index 4fa40b5..6967ee0 100644 --- a/scripts/model_sync.py +++ b/scripts/model_sync.py @@ -52,14 +52,18 @@ def main(): print("[1] Model:", line, "EXIST!") if line.startswith('Nv'): net = line.split('-')[2] - if len(net.split('_')) == 3: - name, seg, ncls = net.split('_') - 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' + if net.startswith("yolov5"): + if len(net.split('_')) == 3: + name, seg, ncls = net.split('_') + 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: - 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' + engine_fn = os.path.splitext(model_file)[0] + '.engine' + online_fn = os.path.splitext(model_file)[0] + '-online.engine' if not os.path.exists(engine_fn) and not os.path.exists(online_fn): if net.startswith("yolov5"): if len(net.split('_')) == 3: