`export.py` return exported files/dirs (#6343)
Browse files* `export.py` return exported files/dirs
* Path to str
export.py
CHANGED
@@ -434,16 +434,17 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
|
|
434 |
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
|
435 |
|
436 |
# Exports
|
|
|
437 |
if 'torchscript' in include:
|
438 |
-
f = export_torchscript(model, im, file, optimize)
|
439 |
if 'engine' in include: # TensorRT required before ONNX
|
440 |
-
f = export_engine(model, im, file, train, half, simplify, workspace, verbose)
|
441 |
if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX
|
442 |
-
f = export_onnx(model, im, file, opset, train, dynamic, simplify)
|
443 |
if 'openvino' in include:
|
444 |
-
f = export_openvino(model, im, file)
|
445 |
if 'coreml' in include:
|
446 |
-
_, f = export_coreml(model, im, file)
|
447 |
|
448 |
# TensorFlow Exports
|
449 |
if any(tf_exports):
|
@@ -451,25 +452,27 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
|
|
451 |
if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
|
452 |
check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
|
453 |
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
|
454 |
-
model, f = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
|
455 |
-
|
456 |
-
|
457 |
if pb or tfjs: # pb prerequisite to tfjs
|
458 |
-
f = export_pb(model, im, file)
|
459 |
if tflite or edgetpu:
|
460 |
-
f = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
|
461 |
if edgetpu:
|
462 |
-
f = export_edgetpu(model, im, file)
|
463 |
if tfjs:
|
464 |
-
f = export_tfjs(model, im, file)
|
465 |
|
466 |
# Finish
|
|
|
467 |
LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
|
468 |
f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
|
469 |
f"\nVisualize with https://netron.app"
|
470 |
-
f"\nDetect with `python detect.py --weights {f}`"
|
471 |
-
f" or `model = torch.hub.load('ultralytics/yolov5', 'custom', '{f}')"
|
472 |
-
f"\nValidate with `python val.py --weights {f}`")
|
|
|
473 |
|
474 |
|
475 |
def parse_opt():
|
|
|
434 |
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
|
435 |
|
436 |
# Exports
|
437 |
+
f = [''] * 10 # exported filenames
|
438 |
if 'torchscript' in include:
|
439 |
+
f[0] = export_torchscript(model, im, file, optimize)
|
440 |
if 'engine' in include: # TensorRT required before ONNX
|
441 |
+
f[1] = export_engine(model, im, file, train, half, simplify, workspace, verbose)
|
442 |
if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX
|
443 |
+
f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
|
444 |
if 'openvino' in include:
|
445 |
+
f[3] = export_openvino(model, im, file)
|
446 |
if 'coreml' in include:
|
447 |
+
_, f[4] = export_coreml(model, im, file)
|
448 |
|
449 |
# TensorFlow Exports
|
450 |
if any(tf_exports):
|
|
|
452 |
if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
|
453 |
check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
|
454 |
assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.'
|
455 |
+
model, f[5] = export_saved_model(model, im, file, dynamic, tf_nms=nms or agnostic_nms or tfjs,
|
456 |
+
agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
|
457 |
+
topk_all=topk_all, conf_thres=conf_thres, iou_thres=iou_thres) # keras model
|
458 |
if pb or tfjs: # pb prerequisite to tfjs
|
459 |
+
f[6] = export_pb(model, im, file)
|
460 |
if tflite or edgetpu:
|
461 |
+
f[7] = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
|
462 |
if edgetpu:
|
463 |
+
f[8] = export_edgetpu(model, im, file)
|
464 |
if tfjs:
|
465 |
+
f[9] = export_tfjs(model, im, file)
|
466 |
|
467 |
# Finish
|
468 |
+
f = [str(x) for x in f if x] # filter out '' and None
|
469 |
LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
|
470 |
f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
|
471 |
f"\nVisualize with https://netron.app"
|
472 |
+
f"\nDetect with `python detect.py --weights {f[-1]}`"
|
473 |
+
f" or `model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}')"
|
474 |
+
f"\nValidate with `python val.py --weights {f[-1]}`")
|
475 |
+
return f # return list of exported files/dirs
|
476 |
|
477 |
|
478 |
def parse_opt():
|