diff --git a/export.py b/export.py index 2ef415c1a..b2b6d5b6f 100644 --- a/export.py +++ b/export.py @@ -2,6 +2,7 @@ import contextlib import json import os +import shutil import platform import re import subprocess @@ -18,36 +19,58 @@ ROOT = FILE.parents[0] # YOLO root directory if str(ROOT) not in sys.path: sys.path.append(str(ROOT)) # add ROOT to PATH -if platform.system() != 'Windows': +if platform.system() != "Windows": ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative from models.experimental import attempt_load, End2End -from models.yolo import ClassificationModel, Detect, DDetect, DualDetect, DualDDetect, DetectionModel, SegmentationModel +from models.yolo import ( + ClassificationModel, + Detect, + DDetect, + DualDetect, + DualDDetect, + DetectionModel, + SegmentationModel, +) from utils.dataloaders import LoadImages -from utils.general import (LOGGER, Profile, check_dataset, check_img_size, check_requirements, check_version, - check_yaml, colorstr, file_size, get_default_args, print_args, url2file, yaml_save) +from utils.general import ( + LOGGER, + Profile, + check_dataset, + check_img_size, + check_requirements, + check_version, + check_yaml, + colorstr, + file_size, + get_default_args, + print_args, + url2file, + yaml_save, +) from utils.torch_utils import select_device, smart_inference_mode -MACOS = platform.system() == 'Darwin' # macOS environment +MACOS = platform.system() == "Darwin" # macOS environment def export_formats(): # YOLO export formats x = [ - ['PyTorch', '-', '.pt', True, True], - ['TorchScript', 'torchscript', '.torchscript', True, True], - ['ONNX', 'onnx', '.onnx', True, True], - ['ONNX END2END', 'onnx_end2end', '_end2end.onnx', True, True], - ['OpenVINO', 'openvino', '_openvino_model', True, False], - ['TensorRT', 'engine', '.engine', False, True], - ['CoreML', 'coreml', '.mlmodel', True, False], - ['TensorFlow SavedModel', 'saved_model', '_saved_model', True, True], - ['TensorFlow GraphDef', 'pb', '.pb', True, True], - ['TensorFlow Lite', 'tflite', '.tflite', True, False], - ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False], - ['TensorFlow.js', 'tfjs', '_web_model', False, False], - ['PaddlePaddle', 'paddle', '_paddle_model', True, True],] - return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU']) + ["PyTorch", "-", ".pt", True, True], + ["TorchScript", "torchscript", ".torchscript", True, True], + ["ONNX", "onnx", ".onnx", True, True], + ["ONNX END2END", "onnx_end2end", "_end2end.onnx", True, True], + ["OpenVINO", "openvino", "_openvino_model", True, False], + ["TensorRT", "engine", ".engine", False, True], + ["CoreML", "coreml", ".mlmodel", True, False], + ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True], + ["TensorFlow GraphDef", "pb", ".pb", True, True], + ["TensorFlow Lite", "tflite", ".tflite", True, False], + ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", False, False], + ["TensorFlow.js", "tfjs", "_web_model", False, False], + ["PaddlePaddle", "paddle", "_paddle_model", True, True], + ] + return pd.DataFrame(x, columns=["Format", "Argument", "Suffix", "CPU", "GPU"]) def try_export(inner_func): @@ -55,52 +78,64 @@ def try_export(inner_func): inner_args = get_default_args(inner_func) def outer_func(*args, **kwargs): - prefix = inner_args['prefix'] + prefix = inner_args["prefix"] try: with Profile() as dt: f, model = inner_func(*args, **kwargs) - LOGGER.info(f'{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)') + LOGGER.info( + f"{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)" + ) return f, model except Exception as e: - LOGGER.info(f'{prefix} export failure ❌ {dt.t:.1f}s: {e}') + LOGGER.info(f"{prefix} export failure ❌ {dt.t:.1f}s: {e}") return None, None return outer_func @try_export -def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')): +def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")): # YOLO TorchScript model export - LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...') - f = file.with_suffix('.torchscript') + LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...") + f = file.with_suffix(".torchscript") ts = torch.jit.trace(model, im, strict=False) d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names} - extra_files = {'config.txt': json.dumps(d)} # torch._C.ExtraFilesMap() + extra_files = {"config.txt": json.dumps(d)} # torch._C.ExtraFilesMap() if optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html - optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files) + optimize_for_mobile(ts)._save_for_lite_interpreter( + str(f), _extra_files=extra_files + ) else: ts.save(str(f), _extra_files=extra_files) return f, None @try_export -def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX:')): +def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")): # YOLO ONNX export - check_requirements('onnx') + check_requirements("onnx") import onnx - LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') - f = file.with_suffix('.onnx') + LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...") + f = file.with_suffix(".onnx") - output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0'] + output_names = ( + ["output0", "output1"] if isinstance(model, SegmentationModel) else ["output0"] + ) if dynamic: - dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640) + dynamic = { + "images": {0: "batch", 2: "height", 3: "width"} + } # shape(1,3,640,640) if isinstance(model, SegmentationModel): - dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) - dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160) + dynamic["output0"] = {0: "batch", 1: "anchors"} # shape(1,25200,85) + dynamic["output1"] = { + 0: "batch", + 2: "mask_height", + 3: "mask_width", + } # shape(1,32,160,160) elif isinstance(model, DetectionModel): - dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85) + dynamic["output0"] = {0: "batch", 1: "anchors"} # shape(1,25200,85) torch.onnx.export( model.cpu() if dynamic else model, # --dynamic only compatible with cpu @@ -109,16 +144,17 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX verbose=False, opset_version=opset, do_constant_folding=True, - input_names=['images'], + input_names=["images"], output_names=output_names, - dynamic_axes=dynamic or None) + dynamic_axes=dynamic or None, + ) # Checks model_onnx = onnx.load(f) # load onnx model onnx.checker.check_model(model_onnx) # check onnx model # Metadata - d = {'stride': int(max(model.stride)), 'names': model.names} + d = {"stride": int(max(model.stride)), "names": model.names} for k, v in d.items(): meta = model_onnx.metadata_props.add() meta.key, meta.value = k, str(v) @@ -128,52 +164,81 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX if simplify: try: cuda = torch.cuda.is_available() - check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1')) + check_requirements( + ("onnxruntime-gpu" if cuda else "onnxruntime", "onnx-simplifier>=0.4.1") + ) import onnxsim - LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') + LOGGER.info( + f"{prefix} simplifying with onnx-simplifier {onnxsim.__version__}..." + ) model_onnx, check = onnxsim.simplify(model_onnx) - assert check, 'assert check failed' + assert check, "assert check failed" onnx.save(model_onnx, f) except Exception as e: - LOGGER.info(f'{prefix} simplifier failure: {e}') + LOGGER.info(f"{prefix} simplifier failure: {e}") return f, model_onnx - + @try_export -def export_onnx_end2end(model, im, file, simplify, topk_all, iou_thres, conf_thres, device, labels, prefix=colorstr('ONNX END2END:')): +def export_onnx_end2end( + model, + im, + file, + simplify, + topk_all, + iou_thres, + conf_thres, + device, + labels, + prefix=colorstr("ONNX END2END:"), +): # YOLO ONNX export - check_requirements('onnx') + check_requirements("onnx") import onnx - LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') + + LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...") f = os.path.splitext(file)[0] + "-end2end.onnx" - batch_size = 'batch' + batch_size = "batch" - dynamic_axes = {'images': {0 : 'batch', 2: 'height', 3:'width'}, } # variable length axes + dynamic_axes = { + "images": {0: "batch", 2: "height", 3: "width"}, + } # variable length axes output_axes = { - 'num_dets': {0: 'batch'}, - 'det_boxes': {0: 'batch'}, - 'det_scores': {0: 'batch'}, - 'det_classes': {0: 'batch'}, - } + "num_dets": {0: "batch"}, + "det_boxes": {0: "batch"}, + "det_scores": {0: "batch"}, + "det_classes": {0: "batch"}, + } dynamic_axes.update(output_axes) - model = End2End(model, topk_all, iou_thres, conf_thres, None ,device, labels) - - output_names = ['num_dets', 'det_boxes', 'det_scores', 'det_classes'] - shapes = [ batch_size, 1, batch_size, topk_all, 4, - batch_size, topk_all, batch_size, topk_all] - - torch.onnx.export(model, - im, - f, - verbose=False, - export_params=True, # store the trained parameter weights inside the model file - opset_version=12, - do_constant_folding=True, # whether to execute constant folding for optimization - input_names=['images'], - output_names=output_names, - dynamic_axes=dynamic_axes) + model = End2End(model, topk_all, iou_thres, conf_thres, None, device, labels) + + output_names = ["num_dets", "det_boxes", "det_scores", "det_classes"] + shapes = [ + batch_size, + 1, + batch_size, + topk_all, + 4, + batch_size, + topk_all, + batch_size, + topk_all, + ] + + torch.onnx.export( + model, + im, + f, + verbose=False, + export_params=True, # store the trained parameter weights inside the model file + opset_version=12, + do_constant_folding=True, # whether to execute constant folding for optimization + input_names=["images"], + output_names=output_names, + dynamic_axes=dynamic_axes, + ) # Checks model_onnx = onnx.load(f) # load onnx model @@ -186,98 +251,128 @@ def export_onnx_end2end(model, im, file, simplify, topk_all, iou_thres, conf_thr try: import onnxsim - print('\nStarting to simplify ONNX...') + print("\nStarting to simplify ONNX...") model_onnx, check = onnxsim.simplify(model_onnx) - assert check, 'assert check failed' + assert check, "assert check failed" except Exception as e: - print(f'Simplifier failure: {e}') + print(f"Simplifier failure: {e}") # print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model - onnx.save(model_onnx,f) - print('ONNX export success, saved as %s' % f) + onnx.save(model_onnx, f) + print("ONNX export success, saved as %s" % f) return f, model_onnx @try_export -def export_openvino(file, metadata, half, prefix=colorstr('OpenVINO:')): +def export_openvino(file, metadata, half, prefix=colorstr("OpenVINO:")): # YOLO OpenVINO export - check_requirements('openvino-dev') # requires openvino-dev: https://pypi.org/project/openvino-dev/ + check_requirements( + "openvino-dev" + ) # requires openvino-dev: https://pypi.org/project/openvino-dev/ import openvino.inference_engine as ie - LOGGER.info(f'\n{prefix} starting export with openvino {ie.__version__}...') - f = str(file).replace('.pt', f'_openvino_model{os.sep}') + LOGGER.info(f"\n{prefix} starting export with openvino {ie.__version__}...") + 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'}" - #cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} {"--compress_to_fp16" if half else ""}" + # cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} --data_type {'FP16' if half else 'FP32'}" + # cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} {"--compress_to_fp16" if half else ""}" half_arg = "--compress_to_fp16" if half else "" cmd = f"mo --input_model {file.with_suffix('.onnx')} --output_dir {f} {half_arg}" subprocess.run(cmd.split(), check=True, env=os.environ) # export - yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata) # add metadata.yaml + yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata) # add metadata.yaml return f, None @try_export -def export_paddle(model, im, file, metadata, prefix=colorstr('PaddlePaddle:')): +def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")): # YOLO Paddle export - check_requirements(('paddlepaddle', 'x2paddle')) + check_requirements(("paddlepaddle", "x2paddle")) import x2paddle from x2paddle.convert import pytorch2paddle - LOGGER.info(f'\n{prefix} starting export with X2Paddle {x2paddle.__version__}...') - f = str(file).replace('.pt', f'_paddle_model{os.sep}') + LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...") + f = str(file).replace(".pt", f"_paddle_model{os.sep}") - pytorch2paddle(module=model, save_dir=f, jit_type='trace', input_examples=[im]) # export - yaml_save(Path(f) / file.with_suffix('.yaml').name, metadata) # add metadata.yaml + pytorch2paddle( + module=model, save_dir=f, jit_type="trace", input_examples=[im] + ) # export + yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata) # add metadata.yaml return f, None @try_export -def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')): +def export_coreml(model, im, file, int8, half, prefix=colorstr("CoreML:")): # YOLO CoreML export - check_requirements('coremltools') - import coremltools as ct - - LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...') - f = file.with_suffix('.mlmodel') + check_requirements("coremltools") + import coremltools as ct # noqa + LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...") + f = file.with_suffix(".mlpackage") + if f.is_dir(): + shutil.rmtree(f) ts = torch.jit.trace(model, im, strict=False) # TorchScript model - ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])]) - bits, mode = (8, 'kmeans_lut') if int8 else (16, 'linear') if half else (32, None) + ct_model = ct.convert( + ts, + inputs=[ct.ImageType("image", shape=im.shape, scale=1 / 255, bias=[0, 0, 0])], + classifier_config=None, + convert_to="mlprogram", + ) + bits, mode = (8, "kmeans_lut") if int8 else (16, "linear") if half else (32, None) if bits < 32: if MACOS: # quantization only supported on macOS with warnings.catch_warnings(): - warnings.filterwarnings("ignore", category=DeprecationWarning) # suppress numpy==1.20 float warning - ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode) + warnings.filterwarnings( + "ignore", category=DeprecationWarning + ) # suppress numpy==1.20 float warning + ct_model = ct.models.neural_network.quantization_utils.quantize_weights( + ct_model, bits, mode + ) else: - print(f'{prefix} quantization only supported on macOS, skipping...') + print(f"{prefix} quantization only supported on macOS, skipping...") ct_model.save(f) return f, ct_model @try_export -def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')): +def export_engine( + model, + im, + file, + half, + dynamic, + simplify, + workspace=4, + verbose=False, + prefix=colorstr("TensorRT:"), +): # YOLO TensorRT export https://developer.nvidia.com/tensorrt - assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`' + assert ( + im.device.type != "cpu" + ), "export running on CPU but must be on GPU, i.e. `python export.py --device 0`" try: import tensorrt as trt except Exception: - if platform.system() == 'Linux': - check_requirements('nvidia-tensorrt', cmds='-U --index-url https://pypi.ngc.nvidia.com') + if platform.system() == "Linux": + check_requirements( + "nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com" + ) import tensorrt as trt - if trt.__version__[0] == '7': # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 + if ( + trt.__version__[0] == "7" + ): # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 grid = model.model[-1].anchor_grid model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid] export_onnx(model, im, file, 12, dynamic, simplify) # opset 12 model.model[-1].anchor_grid = grid else: # TensorRT >= 8 - check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0 + check_version(trt.__version__, "8.0.0", hard=True) # require tensorrt>=8.0.0 export_onnx(model, im, file, 12, dynamic, simplify) # opset 12 - onnx = file.with_suffix('.onnx') + onnx = file.with_suffix(".onnx") - LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...') - assert onnx.exists(), f'failed to export ONNX file: {onnx}' - f = file.with_suffix('.engine') # TensorRT engine file + LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...") + assert onnx.exists(), f"failed to export ONNX file: {onnx}" + f = file.with_suffix(".engine") # TensorRT engine file logger = trt.Logger(trt.Logger.INFO) if verbose: logger.min_severity = trt.Logger.Severity.VERBOSE @@ -287,11 +382,11 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose config.max_workspace_size = workspace * 1 << 30 # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) # fix TRT 8.4 deprecation notice - flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) network = builder.create_network(flag) parser = trt.OnnxParser(network, logger) if not parser.parse_from_file(str(onnx)): - raise RuntimeError(f'failed to load ONNX file: {onnx}') + raise RuntimeError(f"failed to load ONNX file: {onnx}") inputs = [network.get_input(i) for i in range(network.num_inputs)] outputs = [network.get_output(i) for i in range(network.num_outputs)] @@ -302,97 +397,141 @@ def export_engine(model, im, file, half, dynamic, simplify, workspace=4, verbose if dynamic: if im.shape[0] <= 1: - LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument") + LOGGER.warning( + f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument" + ) profile = builder.create_optimization_profile() for inp in inputs: - profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape) + profile.set_shape( + inp.name, + (1, *im.shape[1:]), + (max(1, im.shape[0] // 2), *im.shape[1:]), + im.shape, + ) config.add_optimization_profile(profile) - LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}') + LOGGER.info( + f"{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}" + ) if builder.platform_has_fast_fp16 and half: config.set_flag(trt.BuilderFlag.FP16) - with builder.build_engine(network, config) as engine, open(f, 'wb') as t: + with builder.build_engine(network, config) as engine, open(f, "wb") as t: t.write(engine.serialize()) return f, None @try_export -def export_saved_model(model, - im, - file, - dynamic, - tf_nms=False, - agnostic_nms=False, - topk_per_class=100, - topk_all=100, - iou_thres=0.45, - conf_thres=0.25, - keras=False, - prefix=colorstr('TensorFlow SavedModel:')): +def export_saved_model( + model, + im, + file, + dynamic, + tf_nms=False, + agnostic_nms=False, + topk_per_class=100, + topk_all=100, + iou_thres=0.45, + conf_thres=0.25, + keras=False, + prefix=colorstr("TensorFlow SavedModel:"), +): # YOLO TensorFlow SavedModel export try: import tensorflow as tf except Exception: - check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}") + check_requirements( + f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}" + ) import tensorflow as tf - from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 + from tensorflow.python.framework.convert_to_constants import ( + convert_variables_to_constants_v2, + ) from models.tf import TFModel - LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') - f = str(file).replace('.pt', '_saved_model') + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") + f = str(file).replace(".pt", "_saved_model") batch_size, ch, *imgsz = list(im.shape) # BCHW tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz) im = tf.zeros((batch_size, *imgsz, ch)) # BHWC order for TensorFlow - _ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) - inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size) - outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) + _ = tf_model.predict( + im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres + ) + inputs = tf.keras.Input( + shape=(*imgsz, ch), batch_size=None if dynamic else batch_size + ) + outputs = tf_model.predict( + inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres + ) keras_model = tf.keras.Model(inputs=inputs, outputs=outputs) keras_model.trainable = False keras_model.summary() if keras: - keras_model.save(f, save_format='tf') + keras_model.save(f, save_format="tf") else: spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype) m = tf.function(lambda x: keras_model(x)) # full model m = m.get_concrete_function(spec) frozen_func = convert_variables_to_constants_v2(m) tfm = tf.Module() - tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec]) + tfm.__call__ = tf.function( + lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec] + ) tfm.__call__(im) - tf.saved_model.save(tfm, - f, - options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) if check_version( - tf.__version__, '2.6') else tf.saved_model.SaveOptions()) + tf.saved_model.save( + tfm, + f, + options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) + if check_version(tf.__version__, "2.6") + else tf.saved_model.SaveOptions(), + ) return f, keras_model @try_export -def export_pb(keras_model, file, prefix=colorstr('TensorFlow GraphDef:')): +def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")): # YOLO TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow import tensorflow as tf - from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 + from tensorflow.python.framework.convert_to_constants import ( + convert_variables_to_constants_v2, + ) - LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') - f = file.with_suffix('.pb') + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") + f = file.with_suffix(".pb") m = tf.function(lambda x: keras_model(x)) # full model - m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)) + m = m.get_concrete_function( + tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype) + ) frozen_func = convert_variables_to_constants_v2(m) frozen_func.graph.as_graph_def() - tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False) + tf.io.write_graph( + graph_or_graph_def=frozen_func.graph, + logdir=str(f.parent), + name=f.name, + as_text=False, + ) return f, None @try_export -def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=colorstr('TensorFlow Lite:')): +def export_tflite( + keras_model, + im, + file, + int8, + data, + nms, + agnostic_nms, + prefix=colorstr("TensorFlow Lite:"), +): # YOLOv5 TensorFlow Lite export import tensorflow as tf - LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") batch_size, ch, *imgsz = list(im.shape) # BCHW - f = str(file).replace('.pt', '-fp16.tflite') + f = str(file).replace(".pt", "-fp16.tflite") converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] @@ -400,14 +539,19 @@ def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=c converter.optimizations = [tf.lite.Optimize.DEFAULT] if int8: from models.tf import representative_dataset_gen - dataset = LoadImages(check_dataset(check_yaml(data))['train'], img_size=imgsz, auto=False) - converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100) + + dataset = LoadImages( + check_dataset(check_yaml(data))["train"], img_size=imgsz, auto=False + ) + converter.representative_dataset = lambda: representative_dataset_gen( + dataset, ncalib=100 + ) converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] converter.target_spec.supported_types = [] converter.inference_input_type = tf.uint8 # or tf.int8 converter.inference_output_type = tf.uint8 # or tf.int8 converter.experimental_new_quantizer = True - f = str(file).replace('.pt', '-int8.tflite') + f = str(file).replace(".pt", "-int8.tflite") if nms or agnostic_nms: converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS) @@ -417,24 +561,38 @@ def export_tflite(keras_model, im, file, int8, data, nms, agnostic_nms, prefix=c @try_export -def export_edgetpu(file, prefix=colorstr('Edge TPU:')): +def export_edgetpu(file, prefix=colorstr("Edge TPU:")): # YOLO Edge TPU export https://coral.ai/docs/edgetpu/models-intro/ - cmd = 'edgetpu_compiler --version' - help_url = 'https://coral.ai/docs/edgetpu/compiler/' - assert platform.system() == 'Linux', f'export only supported on Linux. See {help_url}' - if subprocess.run(f'{cmd} >/dev/null', shell=True).returncode != 0: - LOGGER.info(f'\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}') - sudo = subprocess.run('sudo --version >/dev/null', shell=True).returncode == 0 # sudo installed on system + cmd = "edgetpu_compiler --version" + help_url = "https://coral.ai/docs/edgetpu/compiler/" + assert ( + platform.system() == "Linux" + ), f"export only supported on Linux. See {help_url}" + if subprocess.run(f"{cmd} >/dev/null", shell=True).returncode != 0: + LOGGER.info( + f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}" + ) + sudo = ( + subprocess.run("sudo --version >/dev/null", shell=True).returncode == 0 + ) # sudo installed on system for c in ( - 'curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -', - 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list', - 'sudo apt-get update', 'sudo apt-get install edgetpu-compiler'): - subprocess.run(c if sudo else c.replace('sudo ', ''), shell=True, check=True) - ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1] - - LOGGER.info(f'\n{prefix} starting export with Edge TPU compiler {ver}...') - f = str(file).replace('.pt', '-int8_edgetpu.tflite') # Edge TPU model - f_tfl = str(file).replace('.pt', '-int8.tflite') # TFLite model + "curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -", + 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list', + "sudo apt-get update", + "sudo apt-get install edgetpu-compiler", + ): + subprocess.run( + c if sudo else c.replace("sudo ", ""), shell=True, check=True + ) + ver = ( + subprocess.run(cmd, shell=True, capture_output=True, check=True) + .stdout.decode() + .split()[-1] + ) + + LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...") + f = str(file).replace(".pt", "-int8_edgetpu.tflite") # Edge TPU model + f_tfl = str(file).replace(".pt", "-int8.tflite") # TFLite model cmd = f"edgetpu_compiler -s -d -k 10 --out_dir {file.parent} {f_tfl}" subprocess.run(cmd.split(), check=True) @@ -442,30 +600,35 @@ def export_edgetpu(file, prefix=colorstr('Edge TPU:')): @try_export -def export_tfjs(file, prefix=colorstr('TensorFlow.js:')): +def export_tfjs(file, prefix=colorstr("TensorFlow.js:")): # YOLO TensorFlow.js export - check_requirements('tensorflowjs') + check_requirements("tensorflowjs") import tensorflowjs as tfjs - LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...') - f = str(file).replace('.pt', '_web_model') # js dir - f_pb = file.with_suffix('.pb') # *.pb path - f_json = f'{f}/model.json' # *.json path + LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...") + f = str(file).replace(".pt", "_web_model") # js dir + f_pb = file.with_suffix(".pb") # *.pb path + 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}' + 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.split()) json = Path(f_json).read_text() - with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order + with open(f_json, "w") as j: # sort JSON Identity_* in ascending order subst = re.sub( r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}, ' - r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, ' + r'"Identity.?.?": {"name": "Identity.?.?"}}}', + r'{"outputs": {"Identity": {"name": "Identity"}, ' r'"Identity_1": {"name": "Identity_1"}, ' r'"Identity_2": {"name": "Identity_2"}, ' - r'"Identity_3": {"name": "Identity_3"}}}', json) + r'"Identity_3": {"name": "Identity_3"}}}', + json, + ) j.write(subst) return f, None @@ -478,8 +641,8 @@ def add_tflite_metadata(file, metadata, num_outputs): from tflite_support import metadata as _metadata from tflite_support import metadata_schema_py_generated as _metadata_fb - tmp_file = Path('/tmp/meta.txt') - with open(tmp_file, 'w') as meta_f: + tmp_file = Path("/tmp/meta.txt") + with open(tmp_file, "w") as meta_f: meta_f.write(str(metadata)) model_meta = _metadata_fb.ModelMetadataT() @@ -493,7 +656,9 @@ def add_tflite_metadata(file, metadata, num_outputs): model_meta.subgraphMetadata = [subgraph] b = flatbuffers.Builder(0) - b.Finish(model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) + b.Finish( + model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER + ) metadata_buf = b.Output() populator = _metadata.MetadataPopulator.with_model_file(file) @@ -505,53 +670,78 @@ def add_tflite_metadata(file, metadata, num_outputs): @smart_inference_mode() def run( - data=ROOT / 'data/coco.yaml', # 'dataset.yaml path' - weights=ROOT / 'yolo.pt', # weights path - imgsz=(640, 640), # image (height, width) - batch_size=1, # batch size - device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu - include=('torchscript', 'onnx'), # include formats - half=False, # FP16 half-precision export - inplace=False, # set YOLO Detect() inplace=True - keras=False, # use Keras - optimize=False, # TorchScript: optimize for mobile - int8=False, # CoreML/TF INT8 quantization - dynamic=False, # ONNX/TF/TensorRT: dynamic axes - simplify=False, # ONNX: simplify model - opset=12, # ONNX: opset version - verbose=False, # TensorRT: verbose log - workspace=4, # TensorRT: workspace size (GB) - nms=False, # TF: add NMS to model - agnostic_nms=False, # TF: add agnostic NMS to model - topk_per_class=100, # TF.js NMS: topk per class to keep - topk_all=100, # TF.js NMS: topk for all classes to keep - iou_thres=0.45, # TF.js NMS: IoU threshold - conf_thres=0.25, # TF.js NMS: confidence threshold + data=ROOT / "data/coco.yaml", # 'dataset.yaml path' + weights=ROOT / "yolo.pt", # weights path + imgsz=(640, 640), # image (height, width) + batch_size=1, # batch size + device="cpu", # cuda device, i.e. 0 or 0,1,2,3 or cpu + include=("torchscript", "onnx"), # include formats + half=False, # FP16 half-precision export + inplace=False, # set YOLO Detect() inplace=True + keras=False, # use Keras + optimize=False, # TorchScript: optimize for mobile + int8=False, # CoreML/TF INT8 quantization + dynamic=False, # ONNX/TF/TensorRT: dynamic axes + simplify=False, # ONNX: simplify model + opset=12, # ONNX: opset version + verbose=False, # TensorRT: verbose log + workspace=4, # TensorRT: workspace size (GB) + nms=False, # TF: add NMS to model + agnostic_nms=False, # TF: add agnostic NMS to model + topk_per_class=100, # TF.js NMS: topk per class to keep + topk_all=100, # TF.js NMS: topk for all classes to keep + iou_thres=0.45, # TF.js NMS: IoU threshold + conf_thres=0.25, # TF.js NMS: confidence threshold ): t = time.time() include = [x.lower() for x in include] # to lowercase - fmts = tuple(export_formats()['Argument'][1:]) # --include arguments + fmts = tuple(export_formats()["Argument"][1:]) # --include arguments flags = [x in include for x in fmts] - assert sum(flags) == len(include), f'ERROR: Invalid --include {include}, valid --include arguments are {fmts}' - jit, onnx, onnx_end2end, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags # export booleans - file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights) # PyTorch weights + assert sum(flags) == len( + include + ), f"ERROR: Invalid --include {include}, valid --include arguments are {fmts}" + ( + jit, + onnx, + onnx_end2end, + xml, + engine, + coreml, + saved_model, + pb, + tflite, + edgetpu, + tfjs, + paddle, + ) = flags # export booleans + file = Path( + url2file(weights) if str(weights).startswith(("http:/", "https:/")) else weights + ) # PyTorch weights # Load PyTorch model device = select_device(device) if half: - assert device.type != 'cpu' or coreml, '--half only compatible with GPU export, i.e. use --device 0' - assert not dynamic, '--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both' - model = attempt_load(weights, device=device, inplace=True, fuse=True) # load FP32 model + assert ( + device.type != "cpu" or coreml + ), "--half only compatible with GPU export, i.e. use --device 0" + assert not dynamic, "--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both" + model = attempt_load( + weights, device=device, inplace=True, fuse=True + ) # load FP32 model # Checks imgsz *= 2 if len(imgsz) == 1 else 1 # expand if optimize: - assert device.type == 'cpu', '--optimize not compatible with cuda devices, i.e. use --device cpu' + assert ( + device.type == "cpu" + ), "--optimize not compatible with cuda devices, i.e. use --device cpu" # Input gs = int(max(model.stride)) # grid size (max stride) imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples - im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection + im = torch.zeros(batch_size, 3, *imgsz).to( + device + ) # image size(1,3,320,192) BCHW iDetection # Update model model.eval() @@ -565,23 +755,44 @@ def run( y = model(im) # dry runs if half and not coreml: im, model = im.half(), model.half() # to FP16 - shape = tuple((y[0] if isinstance(y, (tuple, list)) else y).shape) # model output shape - metadata = {'stride': int(max(model.stride)), 'names': model.names} # model metadata - LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)") + shape = tuple( + (y[0] if isinstance(y, (tuple, list)) else y).shape + ) # model output shape + metadata = { + "stride": int(max(model.stride)), + "names": model.names, + } # model metadata + LOGGER.info( + f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)" + ) # Exports - f = [''] * len(fmts) # exported filenames - warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning) # suppress TracerWarning + f = [""] * len(fmts) # exported filenames + warnings.filterwarnings( + action="ignore", category=torch.jit.TracerWarning + ) # suppress TracerWarning if jit: # TorchScript f[0], _ = export_torchscript(model, im, file, optimize) if engine: # TensorRT required before ONNX - f[1], _ = export_engine(model, im, file, half, dynamic, simplify, workspace, verbose) + f[1], _ = export_engine( + model, im, file, half, dynamic, simplify, workspace, verbose + ) if onnx or xml: # OpenVINO requires ONNX f[2], _ = export_onnx(model, im, file, opset, dynamic, simplify) if onnx_end2end: if isinstance(model, DetectionModel): labels = model.names - f[2], _ = export_onnx_end2end(model, im, file, simplify, topk_all, iou_thres, conf_thres, device, len(labels)) + f[2], _ = export_onnx_end2end( + model, + im, + file, + simplify, + topk_all, + iou_thres, + conf_thres, + device, + len(labels), + ) else: raise RuntimeError("The model is not a DetectionModel.") if xml: # OpenVINO @@ -589,26 +800,42 @@ def run( if coreml: # CoreML f[4], _ = export_coreml(model, im, file, int8, half) if any((saved_model, pb, tflite, edgetpu, tfjs)): # TensorFlow formats - assert not tflite or not tfjs, 'TFLite and TF.js models must be exported separately, please pass only one type.' - assert not isinstance(model, ClassificationModel), 'ClassificationModel export to TF formats not yet supported.' - f[5], s_model = export_saved_model(model.cpu(), - im, - file, - dynamic, - tf_nms=nms or agnostic_nms or tfjs, - agnostic_nms=agnostic_nms or tfjs, - topk_per_class=topk_per_class, - topk_all=topk_all, - iou_thres=iou_thres, - conf_thres=conf_thres, - keras=keras) + assert ( + not tflite or not tfjs + ), "TFLite and TF.js models must be exported separately, please pass only one type." + assert not isinstance( + model, ClassificationModel + ), "ClassificationModel export to TF formats not yet supported." + f[5], s_model = export_saved_model( + model.cpu(), + im, + file, + dynamic, + tf_nms=nms or agnostic_nms or tfjs, + agnostic_nms=agnostic_nms or tfjs, + topk_per_class=topk_per_class, + topk_all=topk_all, + iou_thres=iou_thres, + conf_thres=conf_thres, + keras=keras, + ) if pb or tfjs: # pb prerequisite to tfjs f[6], _ = export_pb(s_model, file) if tflite or edgetpu: - f[7], _ = export_tflite(s_model, im, file, int8 or edgetpu, data=data, nms=nms, agnostic_nms=agnostic_nms) + f[7], _ = export_tflite( + s_model, + im, + file, + int8 or edgetpu, + data=data, + nms=nms, + agnostic_nms=agnostic_nms, + ) if edgetpu: f[8], _ = export_edgetpu(file) - add_tflite_metadata(f[8] or f[7], metadata, num_outputs=len(s_model.outputs)) + add_tflite_metadata( + f[8] or f[7], metadata, num_outputs=len(s_model.outputs) + ) if tfjs: f[9], _ = export_tfjs(file) if paddle: # PaddlePaddle @@ -617,56 +844,138 @@ def run( # Finish f = [str(x) for x in f if x] # filter out '' and None if any(f): - cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel)) # type - dir = Path('segment' if seg else 'classify' if cls else '') - h = '--half' if half else '' # --half FP16 inference arg - s = "# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference" if cls else \ - "# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference" if seg else '' + cls, det, seg = ( + isinstance(model, x) + for x in (ClassificationModel, DetectionModel, SegmentationModel) + ) # type + dir = Path("segment" if seg else "classify" if cls else "") + h = "--half" if half else "" # --half FP16 inference arg + s = ( + "# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference" + if cls + else "# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference" + if seg + else "" + ) if onnx_end2end: - LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)' - f"\nResults saved to {colorstr('bold', file.parent.resolve())}" - f"\nVisualize: https://netron.app") + LOGGER.info( + f'\nExport complete ({time.time() - t:.1f}s)' + f"\nResults saved to {colorstr('bold', file.parent.resolve())}" + f"\nVisualize: https://netron.app" + ) else: - LOGGER.info(f'\nExport complete ({time.time() - t:.1f}s)' - f"\nResults saved to {colorstr('bold', file.parent.resolve())}" - f"\nDetect: python {dir / ('detect.py' if det else 'predict.py')} --weights {f[-1]} {h}" - f"\nValidate: python {dir / 'val.py'} --weights {f[-1]} {h}" - f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}') {s}" - f"\nVisualize: https://netron.app") + LOGGER.info( + f'\nExport complete ({time.time() - t:.1f}s)' + f"\nResults saved to {colorstr('bold', file.parent.resolve())}" + f"\nDetect: python {dir / ('detect.py' if det else 'predict.py')} --weights {f[-1]} {h}" + f"\nValidate: python {dir / 'val.py'} --weights {f[-1]} {h}" + f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}') {s}" + f"\nVisualize: https://netron.app" + ) return f # return list of exported files/dirs +class IOSDetectModel(torch.nn.Module): + """Wrap an Ultralytics YOLO model for Apple iOS CoreML export.""" + + def __init__(self, model, im): + """Initialize the IOSDetectModel class with a YOLO model and example image.""" + super().__init__() + _, _, h, w = im.shape # batch, channel, height, width + self.model = model + self.nc = len(model.names) # number of classes + if w == h: + self.normalize = 1.0 / w # scalar + else: + self.normalize = torch.tensor([1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h]) # broadcast (slower, smaller) + def forward(self, x): + """Normalize predictions of object detection model with input size-dependent factors.""" + xywh, cls = self.model(x)[0].transpose(0, 1).split((4, self.nc), 1) + return cls, xywh * self.normalize # confidence (3780, 80), coordinates (3780, 4) def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--data', type=str, default=ROOT / 'data/coco.yaml', help='dataset.yaml path') - parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolo.pt', help='model.pt path(s)') - parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)') - parser.add_argument('--batch-size', type=int, default=1, help='batch size') - parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') - parser.add_argument('--half', action='store_true', help='FP16 half-precision export') - parser.add_argument('--inplace', action='store_true', help='set YOLO Detect() inplace=True') - parser.add_argument('--keras', action='store_true', help='TF: use Keras') - parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile') - parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization') - parser.add_argument('--dynamic', action='store_true', help='ONNX/TF/TensorRT: dynamic axes') - parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') - parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version') - parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log') - parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)') - parser.add_argument('--nms', action='store_true', help='TF: add NMS to model') - parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model') - parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep') - parser.add_argument('--topk-all', type=int, default=100, help='ONNX END2END/TF.js NMS: topk for all classes to keep') - parser.add_argument('--iou-thres', type=float, default=0.45, help='ONNX END2END/TF.js NMS: IoU threshold') - parser.add_argument('--conf-thres', type=float, default=0.25, help='ONNX END2END/TF.js NMS: confidence threshold') parser.add_argument( - '--include', - nargs='+', - default=['torchscript'], - help='torchscript, onnx, onnx_end2end, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle') + "--data", type=str, default=ROOT / "data/coco.yaml", help="dataset.yaml path" + ) + parser.add_argument( + "--weights", + nargs="+", + type=str, + default=ROOT / "yolo.pt", + help="model.pt path(s)", + ) + parser.add_argument( + "--imgsz", + "--img", + "--img-size", + nargs="+", + type=int, + default=[640, 640], + help="image (h, w)", + ) + parser.add_argument("--batch-size", type=int, default=1, help="batch size") + parser.add_argument( + "--device", default="cpu", help="cuda device, i.e. 0 or 0,1,2,3 or cpu" + ) + parser.add_argument( + "--half", action="store_true", help="FP16 half-precision export" + ) + parser.add_argument( + "--inplace", action="store_true", help="set YOLO Detect() inplace=True" + ) + parser.add_argument("--keras", action="store_true", help="TF: use Keras") + parser.add_argument( + "--optimize", action="store_true", help="TorchScript: optimize for mobile" + ) + parser.add_argument( + "--int8", action="store_true", help="CoreML/TF INT8 quantization" + ) + parser.add_argument( + "--dynamic", action="store_true", help="ONNX/TF/TensorRT: dynamic axes" + ) + parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model") + parser.add_argument("--opset", type=int, default=12, help="ONNX: opset version") + parser.add_argument("--verbose", action="store_true", help="TensorRT: verbose log") + parser.add_argument( + "--workspace", type=int, default=4, help="TensorRT: workspace size (GB)" + ) + parser.add_argument("--nms", action="store_true", help="TF: add NMS to model") + parser.add_argument( + "--agnostic-nms", action="store_true", help="TF: add agnostic NMS to model" + ) + parser.add_argument( + "--topk-per-class", + type=int, + default=100, + help="TF.js NMS: topk per class to keep", + ) + parser.add_argument( + "--topk-all", + type=int, + default=100, + help="ONNX END2END/TF.js NMS: topk for all classes to keep", + ) + parser.add_argument( + "--iou-thres", + type=float, + default=0.45, + help="ONNX END2END/TF.js NMS: IoU threshold", + ) + parser.add_argument( + "--conf-thres", + type=float, + default=0.25, + help="ONNX END2END/TF.js NMS: confidence threshold", + ) + parser.add_argument( + "--include", + nargs="+", + default=["torchscript"], + help="torchscript, onnx, onnx_end2end, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle", + ) opt = parser.parse_args() - if 'onnx_end2end' in opt.include: + if "onnx_end2end" in opt.include: opt.simplify = True opt.dynamic = True opt.inplace = True @@ -677,7 +986,7 @@ def parse_opt(): def main(opt): - for opt.weights in (opt.weights if isinstance(opt.weights, list) else [opt.weights]): + for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]: run(**vars(opt))