From e5a0573ddedaafd1877fb5ae732fe5079f586b9a Mon Sep 17 00:00:00 2001 From: wang-xinyu Date: Tue, 6 Sep 2022 19:52:31 +0800 Subject: [PATCH] add gen wts choices --- yolov5/gen_wts.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/yolov5/gen_wts.py b/yolov5/gen_wts.py index c20fa50..d55ad9e 100644 --- a/yolov5/gen_wts.py +++ b/yolov5/gen_wts.py @@ -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')