diff --git a/README.md b/README.md index e4a98215..95f1be46 100644 --- a/README.md +++ b/README.md @@ -299,7 +299,7 @@ Video speed is adjusted approximately 50 times slower than actual speed. docker run --rm -it \ -v `pwd`:/workdir \ -w /workdir \ - ghcr.io/pinto0309/onnx2tf:1.26.4 + ghcr.io/pinto0309/onnx2tf:1.26.5 or @@ -307,7 +307,7 @@ Video speed is adjusted approximately 50 times slower than actual speed. docker run --rm -it \ -v `pwd`:/workdir \ -w /workdir \ - docker.io/pinto0309/onnx2tf:1.26.4 + docker.io/pinto0309/onnx2tf:1.26.5 or @@ -1526,6 +1526,7 @@ usage: onnx2tf [-ois OVERWRITE_INPUT_SHAPE [OVERWRITE_INPUT_SHAPE ...]] [-nlt] [-onwdt] +[-snms {v4,v5}] [-k KEEP_NCW_OR_NCHW_OR_NCDHW_INPUT_NAMES [KEEP_NCW_OR_NCHW_OR_NCDHW_INPUT_NAMES ...]] [-kt KEEP_NWC_OR_NHWC_OR_NDHWC_INPUT_NAMES [KEEP_NWC_OR_NHWC_OR_NDHWC_INPUT_NAMES ...]] [-kat KEEP_SHAPE_ABSOLUTELY_INPUT_NAMES [KEEP_SHAPE_ABSOLUTELY_INPUT_NAMES ...]] @@ -1725,6 +1726,12 @@ optional arguments: enable --output_nms_with_dynamic_tensor: output_tensor_shape: [N, 7] + -snms {v4,v5}, --switch_nms_version {v4,v5} + Switch the NMS version to V4 or V5 to convert. + e.g. + NonMaxSuppressionV4(default): --switch_nms_version v4 + NonMaxSuppressionV5: --switch_nms_version v5 + -k KEEP_NCW_OR_NCHW_OR_NCDHW_INPUT_NAMES [KEEP_NCW_OR_NCHW_OR_NCDHW_INPUT_NAMES ...], \ --keep_ncw_or_nchw_or_ncdhw_input_names KEEP_NCW_OR_NCHW_OR_NCDHW_INPUT_NAMES \ [KEEP_NCW_OR_NCHW_OR_NCDHW_INPUT_NAMES ...] @@ -2010,6 +2017,7 @@ convert( overwrite_input_shape: Union[List[str], NoneType] = None, no_large_tensor: Optional[bool] = False, output_nms_with_dynamic_tensor: Optional[bool] = False, + switch_nms_version: Optional[str] = 'v4', keep_ncw_or_nchw_or_ncdhw_input_names: Union[List[str], NoneType] = None, keep_nwc_or_nhwc_or_ndhwc_input_names: Union[List[str], NoneType] = None, keep_shape_absolutely_input_names: Optional[List[str]] = None, @@ -2215,6 +2223,12 @@ convert( enable --output_nms_with_dynamic_tensor: output_tensor_shape: [N, 7] + switch_nms_version {v4,v5} + Switch the NMS version to V4 or V5 to convert. + e.g. + NonMaxSuppressionV4(default): switch_nms_version="v4" + NonMaxSuppressionV5: switch_nms_version="v5" + keep_ncw_or_nchw_or_ncdhw_input_names: Optional[List[str]] Holds the NCW or NCHW or NCDHW of the input shape for the specified INPUT OP names. If a nonexistent INPUT OP name is specified, it is ignored. diff --git a/onnx2tf/__init__.py b/onnx2tf/__init__.py index a5bb143d..135141c0 100644 --- a/onnx2tf/__init__.py +++ b/onnx2tf/__init__.py @@ -1,3 +1,3 @@ from onnx2tf.onnx2tf import convert, main -__version__ = '1.26.4' +__version__ = '1.26.5' diff --git a/onnx2tf/onnx2tf.py b/onnx2tf/onnx2tf.py index b58d570e..5162c9b2 100644 --- a/onnx2tf/onnx2tf.py +++ b/onnx2tf/onnx2tf.py @@ -80,6 +80,7 @@ def convert( overwrite_input_shape: Optional[List[str]] = None, no_large_tensor: Optional[bool] = False, output_nms_with_dynamic_tensor: Optional[bool] = False, + switch_nms_version: Optional[str] = 'v4', keep_ncw_or_nchw_or_ncdhw_input_names: Optional[List[str]] = None, keep_nwc_or_nhwc_or_ndhwc_input_names: Optional[List[str]] = None, keep_shape_absolutely_input_names: Optional[List[str]] = None, @@ -270,6 +271,12 @@ def convert( enable --output_nms_with_dynamic_tensor:\n output_tensor_shape: [N, 7] + switch_nms_version: Optional[str] + Switch the NMS version to V4 or V5 to convert.\n\n + e.g.\n + NonMaxSuppressionV4(default): --switch_nms_version v4\n + NonMaxSuppressionV5: --switch_nms_version v5 + keep_ncw_or_nchw_or_ncdhw_input_names: Optional[List[str]] Holds the NCW or NCHW or NCDHW of the input shape for the specified INPUT OP names.\n If a nonexistent INPUT OP name is specified, it is ignored.\n @@ -921,6 +928,7 @@ def sanitizing(node): 'mvn_epsilon': mvn_epsilon, 'output_signaturedefs': output_signaturedefs, 'output_nms_with_dynamic_tensor': output_nms_with_dynamic_tensor, + 'switch_nms_version': switch_nms_version, 'output_integer_quantized_tflite': output_integer_quantized_tflite, 'gelu_replace_op_names': {}, 'space_to_depth_replace_op_names': {}, @@ -2233,6 +2241,18 @@ def main(): 'enable --output_nms_with_dynamic_tensor: \n' + ' output_tensor_shape: [N, 7]' ) + parser.add_argument( + '-snms', + '--switch_nms_version', + type=str, + choices=['v4', 'v5'], + default='v4', + help=\ + 'Switch the NMS version to V4 or V5 to convert. \n' + + 'e.g. \n' + + 'NonMaxSuppressionV4(default): --switch_nms_version v4 \n' + + 'NonMaxSuppressionV5: --switch_nms_version v5' + ) parser.add_argument( '-k', '--keep_ncw_or_nchw_or_ncdhw_input_names', @@ -2623,6 +2643,7 @@ def main(): overwrite_input_shape=args.overwrite_input_shape, no_large_tensor=args.no_large_tensor, output_nms_with_dynamic_tensor=args.output_nms_with_dynamic_tensor, + switch_nms_version=args.switch_nms_version, keep_ncw_or_nchw_or_ncdhw_input_names=args.keep_ncw_or_nchw_or_ncdhw_input_names, keep_nwc_or_nhwc_or_ndhwc_input_names=args.keep_nwc_or_nhwc_or_ndhwc_input_names, keep_shape_absolutely_input_names=args.keep_shape_absolutely_input_names, diff --git a/onnx2tf/ops/NonMaxSuppression.py b/onnx2tf/ops/NonMaxSuppression.py index 9d70bac4..1b88dd93 100644 --- a/onnx2tf/ops/NonMaxSuppression.py +++ b/onnx2tf/ops/NonMaxSuppression.py @@ -25,8 +25,9 @@ class NMSLayer(tf_keras.layers.Layer): - def __init__(self): + def __init__(self, switch_nms_version='v4'): super(NMSLayer, self).__init__() + self.switch_nms_version = switch_nms_version @dispatch.add_dispatch_support def non_max_suppression( @@ -40,29 +41,56 @@ def non_max_suppression( name=None, ): with ops.name_scope(name, 'non_max_suppression'): - selected_indices, num_valid = gen_image_ops.non_max_suppression_v4( - boxes=boxes, - scores=scores, - max_output_size=max_output_size \ - if not isinstance(max_output_size, np.ndarray) \ - else tf.convert_to_tensor( - value=max_output_size, - name='max_output_size' - ), - iou_threshold=iou_threshold \ - if not isinstance(iou_threshold, np.ndarray) \ - else tf.convert_to_tensor( - value=iou_threshold, - name='iou_threshold', - ), - score_threshold=score_threshold \ - if not isinstance(score_threshold, np.ndarray) \ - else tf.convert_to_tensor( - value=score_threshold, - name='score_threshold', - ), - pad_to_max_output_size=pad_to_max_output_size, - ) + if self.switch_nms_version == 'v4': + selected_indices, num_valid = gen_image_ops.non_max_suppression_v4( + boxes=boxes, + scores=scores, + max_output_size=max_output_size \ + if not isinstance(max_output_size, np.ndarray) \ + else tf.convert_to_tensor( + value=max_output_size, + name='max_output_size' + ), + iou_threshold=iou_threshold \ + if not isinstance(iou_threshold, np.ndarray) \ + else tf.convert_to_tensor( + value=iou_threshold, + name='iou_threshold', + ), + score_threshold=score_threshold \ + if not isinstance(score_threshold, np.ndarray) \ + else tf.convert_to_tensor( + value=score_threshold, + name='score_threshold', + ), + pad_to_max_output_size=pad_to_max_output_size, + ) + + elif self.switch_nms_version == 'v5': + selected_indices, selected_scores, num_valid = gen_image_ops.non_max_suppression_v5( + boxes=boxes, + scores=scores, + max_output_size=max_output_size \ + if not isinstance(max_output_size, np.ndarray) \ + else tf.convert_to_tensor( + value=max_output_size, + name='max_output_size' + ), + iou_threshold=iou_threshold \ + if not isinstance(iou_threshold, np.ndarray) \ + else tf.convert_to_tensor( + value=iou_threshold, + name='iou_threshold', + ), + score_threshold=score_threshold \ + if not isinstance(score_threshold, np.ndarray) \ + else tf.convert_to_tensor( + value=score_threshold, + name='score_threshold', + ), + soft_nms_sigma=0.0, + pad_to_max_output_size=pad_to_max_output_size, + ) if pad_to_max_output_size: return selected_indices @@ -130,8 +158,8 @@ def make_node( scores = tf_layers_dict[graph_node_input_2.name]['tf_node'] \ if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2 - output_nms_with_dynamic_tensor: bool = \ - kwargs['output_nms_with_dynamic_tensor'] + output_nms_with_dynamic_tensor: bool = kwargs['output_nms_with_dynamic_tensor'] + switch_nms_version: str = kwargs['switch_nms_version'] # Pre-process transpose boxes = pre_process_transpose( @@ -339,7 +367,7 @@ def make_node( axis=0, ) # get the selected boxes indices - nms = NMSLayer() + nms = NMSLayer(switch_nms_version=switch_nms_version) selected_indices = nms( boxes=tf_boxes, scores=tf_scores,