From 064365d8683fd002e9ad789c1e91fa3d021b44f0 Mon Sep 17 00:00:00 2001 From: Johan Bergman <35481994+duran67@users.noreply.github.com> Date: Fri, 20 Jan 2023 23:49:43 +0100 Subject: [PATCH] Update parse_opt() in export.py to work as in train.py (#10789) Update parse_opt() to work as in train.py Change parse_opt() be able to use parse_known_args(), same as in train.py, so export.main() can be called from other script without error. e.g.: from yolov5 import export opt = export.parse_opt(True) opt.weights = opt.include = ("torchscript", "onnx") opt.data = opt.imgsz = [, ] export.main(opt) Signed-off-by: Johan Bergman <35481994+duran67@users.noreply.github.com> Signed-off-by: Johan Bergman <35481994+duran67@users.noreply.github.com> --- export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/export.py b/export.py index 7910178..9ca3441 100644 --- a/export.py +++ b/export.py @@ -610,7 +610,7 @@ def run( return f # return list of exported files/dirs -def parse_opt(): +def parse_opt(known=False): parser = argparse.ArgumentParser() parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)') @@ -638,7 +638,7 @@ def parse_opt(): nargs='+', default=['torchscript'], help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle') - opt = parser.parse_args() + opt = parser.parse_known_args()[0] if known else parser.parse_args() print_args(vars(opt)) return opt