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

Fix gen_wts_yolor

Browse Source
This commit is contained in:
Marcos Luciano
2022-05-28 22:41:48 -03:00
parent fcc2463a92
commit 8ff01e850d
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
utils/gen_wts_yolor.py
View File

@@ -20,7 +20,7 @@ def parse_args():
pt_file, cfg_file = parse_args() pt_file, cfg_file = parse_args()
wts_file = cfg_file.split(".cfg")[0] + ".wts" wts_file = "%s.wts" % cfg_file.rsplit("/")[1].split(".cfg")[0]
device = select_device("cpu") device = select_device("cpu")
model = Darknet(cfg_file).to(device) model = Darknet(cfg_file).to(device)
@@ -41,3 +41,5 @@ with open(wts_file, "w") as f:
conv_count += 1 conv_count += 1
f.write("{}\n".format(conv_count)) f.write("{}\n".format(conv_count))
f.write(wts_write) f.write(wts_write)
os.system("cp %s ./" % cfg_file)