haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add gen wts choices

Browse Source
This commit is contained in:
wang-xinyu 2022-09-06 19:52:31 +08:00
parent fdd6a76ec2
commit e5a0573dde
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
yolov5/gen_wts.py
View File

@ -13,7 +13,8 @@ def parse_args():
parser.add_argument(
'-o', '--output', help='Output (.wts) file path (optional)')
parser.add_argument(
'-t', '--type', type=str, default='', help='determines the model is detection/classification')
'-t', '--type', type=str, default='detect', choices=['detect', 'cls'],
help='determines the model is detection/classification')
args = parser.parse_args()
if not os.path.isfile(args.weights):
raise SystemExit('Invalid input file')