glenn-jocher commited on
Commit
0cf932b
·
unverified ·
1 Parent(s): e2e95b2

`export.py` return exported files/dirs (#6343)

Browse files

* `export.py` return exported files/dirs

* Path to str

Files changed (1) hide show
  1. export.py +18 -15
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
- 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
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():