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>
This commit is contained in:
1 parent
c442a2e993
commit
064365d868
1 file changed
+2
-2
@@ -610,7 +610,7 @@ def run(
|
|||||||
return f # return list of exported files/dirs
|
return f # return list of exported files/dirs
|
||||||
|
|
||||||
|
|
||||||
def parse_opt():
|
def parse_opt(known=False):
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
|
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)')
|
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
|
||||||
@@ -638,7 +638,7 @@ def parse_opt():
|
|||||||
nargs='+',
|
nargs='+',
|
||||||
default=['torchscript'],
|
default=['torchscript'],
|
||||||
help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle')
|
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))
|
print_args(vars(opt))
|
||||||
return opt
|
return opt
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user