diff --git a/export.py b/export.py index 3325598bf489..4a8b50d2d74c 100644 --- a/export.py +++ b/export.py @@ -172,7 +172,7 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst LOGGER.info(f'{prefix} export failure: {e}') -def export_openvino(model, im, file, half, prefix=colorstr('OpenVINO:')): +def export_openvino(file, half, prefix=colorstr('OpenVINO:')): # YOLOv5 OpenVINO export try: check_requirements(('openvino-dev',)) # requires openvino-dev: https://pypi.org/project/openvino-dev/ @@ -182,7 +182,7 @@ def export_openvino(model, im, file, half, prefix=colorstr('OpenVINO:')): f = str(file).replace('.pt', f'_openvino_model{os.sep}') cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}" - subprocess.check_output(cmd, shell=True) + subprocess.check_output(cmd.split()) LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') return f @@ -328,7 +328,7 @@ def export_saved_model(model, return None, None -def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')): +def export_pb(keras_model, file, prefix=colorstr('TensorFlow GraphDef:')): # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow try: import tensorflow as tf @@ -383,7 +383,7 @@ def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=c LOGGER.info(f'\n{prefix} export failure: {e}') -def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): +def export_edgetpu(file, prefix=colorstr('Edge TPU:')): # YOLOv5 Edge TPU export https://coral.ai/docs/edgetpu/models-intro/ try: cmd = 'edgetpu_compiler --version' @@ -404,7 +404,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model cmd = f"edgetpu_compiler -s -o {file.parent} {f_tfl}" - subprocess.run(cmd, shell=True, check=True) + subprocess.run(cmd.split(), check=True) LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') return f @@ -412,7 +412,7 @@ def export_edgetpu(keras_model, im, file, prefix=colorstr('Edge TPU:')): LOGGER.info(f'\n{prefix} export failure: {e}') -def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')): +def export_tfjs(file, prefix=colorstr('TensorFlow.js:')): # YOLOv5 TensorFlow.js export try: check_requirements(('tensorflowjs',)) @@ -426,8 +426,8 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')): f_json = f'{f}/model.json' # *.json path cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \ - f'--output_node_names="Identity,Identity_1,Identity_2,Identity_3" {f_pb} {f}' - subprocess.run(cmd, shell=True) + f'--output_node_names=Identity,Identity_1,Identity_2,Identity_3 {f_pb} {f}' + subprocess.run(cmd.split()) with open(f_json) as j: json = j.read() @@ -524,7 +524,7 @@ def run( print(im.shape[3]) f[2] = export_onnx(model, im, file, opset, train, dynamic, simplify) if xml: # OpenVINO - f[3] = export_openvino(model, im, file, half) + f[3] = export_openvino(file, half) if coreml: _, f[4] = export_coreml(model, im, file, int8, half) @@ -544,13 +544,13 @@ def run( conf_thres=conf_thres, iou_thres=iou_thres) # keras model if pb or tfjs: # pb prerequisite to tfjs - f[6] = export_pb(model, im, file) + f[6] = export_pb(model, file) if tflite or edgetpu: f[7] = export_tflite(model, im, file, int8=int8 or edgetpu, data=data, nms=nms, agnostic_nms=agnostic_nms) if edgetpu: - f[8] = export_edgetpu(model, im, file) + f[8] = export_edgetpu(file) if tfjs: - f[9] = export_tfjs(model, im, file) + f[9] = export_tfjs(file) # Finish f = [str(x) for x in f if x] # filter out '' and None