mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

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 = <model_path>
    opt.include = ("torchscript", "onnx")
    opt.data = <data>
    opt.imgsz = [<height>, <width>]
    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>
pull/10793/head^2
Johan Bergman 2023-01-20 23:49:43 +01:00 committed by GitHub
parent c442a2e993
commit 064365d868
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
export.py
View File

@ -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