Skip to content

Commit cb36846

Browse files
authored
export.py return exported files/dirs (ultralytics#6343)
* `export.py` return exported files/dirs * Path to str
1 parent 95c28a9 commit cb36846

File tree

1 file changed

+18
-15
lines changed

1 file changed

+18
-15
lines changed

export.py

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -434,42 +434,45 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
434434
LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)")
435435

436436
# Exports
437+
f = [''] * 10 # exported filenames
437438
if 'torchscript' in include:
438-
f = export_torchscript(model, im, file, optimize)
439+
f[0] = export_torchscript(model, im, file, optimize)
439440
if 'engine' in include: # TensorRT required before ONNX
440-
f = export_engine(model, im, file, train, half, simplify, workspace, verbose)
441+
f[1] = export_engine(model, im, file, train, half, simplify, workspace, verbose)
441442
if ('onnx' in include) or ('openvino' in include): # OpenVINO requires ONNX
442-
f = export_onnx(model, im, file, opset, train, dynamic, simplify)
443+
f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify)
443444
if 'openvino' in include:
444-
f = export_openvino(model, im, file)
445+
f[3] = export_openvino(model, im, file)
445446
if 'coreml' in include:
446-
_, f = export_coreml(model, im, file)
447+
_, f[4] = export_coreml(model, im, file)
447448

448449
# TensorFlow Exports
449450
if any(tf_exports):
450451
pb, tflite, edgetpu, tfjs = tf_exports[1:]
451452
if int8 or edgetpu: # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
452453
check_requirements(('flatbuffers==1.12',)) # required before `import tensorflow`
453454
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-
agnostic_nms=agnostic_nms or tfjs, topk_per_class=topk_per_class,
456-
topk_all=topk_all, conf_thres=conf_thres, iou_thres=iou_thres) # keras model
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
457458
if pb or tfjs: # pb prerequisite to tfjs
458-
f = export_pb(model, im, file)
459+
f[6] = export_pb(model, im, file)
459460
if tflite or edgetpu:
460-
f = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
461+
f[7] = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, ncalib=100)
461462
if edgetpu:
462-
f = export_edgetpu(model, im, file)
463+
f[8] = export_edgetpu(model, im, file)
463464
if tfjs:
464-
f = export_tfjs(model, im, file)
465+
f[9] = export_tfjs(model, im, file)
465466

466467
# Finish
468+
f = [str(x) for x in f if x] # filter out '' and None
467469
LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)'
468470
f"\nResults saved to {colorstr('bold', file.parent.resolve())}"
469471
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}`")
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
473476

474477

475478
def parse_opt():

0 commit comments

Comments
 (0)