return onnx.ModelProto

Author: PINTO0309

README.md

@@ -35,27 +35,31 @@ $ cd /workdir
 $ sne4onnx -h
 
 usage:
-    sne4onnx [-h] \
-    --input_onnx_file_path INPUT_ONNX_FILE_PATH \
-    --output_onnx_file_path OUTPUT_ONNX_FILE_PATH \
-    --input_op_names INPUT_OP_NAMES \
+    sne4onnx [-h]
+    --input_onnx_file_path INPUT_ONNX_FILE_PATH
+    --input_op_names INPUT_OP_NAMES
     --output_op_names OUTPUT_OP_NAMES
+    [--output_onnx_file_path OUTPUT_ONNX_FILE_PATH]
 
 optional arguments:
   -h, --help
         show this help message and exit
+
   --input_onnx_file_path INPUT_ONNX_FILE_PATH
         Input onnx file path.
-  --output_onnx_file_path OUTPUT_ONNX_FILE_PATH
-        Output onnx file path.
+
   --input_op_names INPUT_OP_NAMES
         List of OP names to specify for the input layer of the model.
         Specify the name of the OP, separated by commas.
         e.g. --input_op_names aaa,bbb,ccc
+
   --output_op_names OUTPUT_OP_NAMES
         List of OP names to specify for the output layer of the model.
         Specify the name of the OP, separated by commas.
         e.g. --output_op_names ddd,eee,fff
+
+  --output_onnx_file_path OUTPUT_ONNX_FILE_PATH
+        Output onnx file path. If not specified, extracted.onnx is output.
 ```
 
 ## 3. In-script Usage
@@ -68,19 +72,16 @@ Help on function extraction in module sne4onnx.onnx_network_extraction:
 
 extraction(
     input_onnx_file_path: str,
-    output_onnx_file_path: str,
     input_op_names: List[str],
-    output_op_names: List[str]
-)
+    output_op_names: List[str],
+    output_onnx_file_path: Union[str, NoneType] = ''
+) -> onnx.onnx_ml_pb2.ModelProto
 
     Parameters
     ----------
     input_onnx_file_path: str
         Input onnx file path.
 
-    output_onnx_file_path: str
-        Output onnx file path.
-
     input_op_names: List[str]
         List of OP names to specify for the input layer of the model.
         Specify the name of the OP, separated by commas.
@@ -90,26 +91,36 @@ extraction(
         List of OP names to specify for the output layer of the model.
         Specify the name of the OP, separated by commas.
         e.g. ['ddd','eee','fff']
+
+    output_onnx_file_path: Optional[str]
+        Output onnx file path.
+        If not specified, .onnx is not output.
+        Default: ''
+
+    Returns
+    -------
+    extracted_graph: onnx.ModelProto
+        Extracted onnx ModelProto
 ```
 
 ## 4. CLI Execution
 ```bash
 $ sne4onnx \
 --input_onnx_file_path input.onnx \
---output_onnx_file_path output.onnx \
 --input_op_names aaa,bbb,ccc \
---output_op_names ddd,eee,fff
+--output_op_names ddd,eee,fff \
+--output_onnx_file_path output.onnx
 ```
 
 ## 5. In-script Execution
 ```python
 from sne4onnx import extraction
 
-extraction(
+extracted_graph = extraction(
   input_onnx_file_path='input.onnx',
-  output_onnx_file_path='output.onnx',
   input_op_names=['aaa', 'bbb', 'ccc'],
   output_op_names=['ddd', 'eee', 'fff'],
+  output_onnx_file_path='output.onnx',
 )
 ```
 
@@ -123,9 +134,9 @@ extraction(
 ```bash
 $ sne4onnx \
 --input_onnx_file_path hitnet_sf_finalpass_720x1280.onnx \
---output_onnx_file_path hitnet_sf_finalpass_720x960_head.onnx \
 --input_op_names 0,1 \
---output_op_names 497,785
+--output_op_names 497,785 \
+--output_onnx_file_path hitnet_sf_finalpass_720x960_head.onnx
 ```
 
 ### 6-3. Extracted

sne4onnx/__init__.py

@@ -1,3 +1,3 @@
 from sne4onnx.onnx_network_extraction import extraction, main
 
-__version__ = '1.0.2'
+__version__ = '1.0.3'

sne4onnx/onnx_network_extraction.py

@@ -1,8 +1,9 @@
 #! /usr/bin/env python
 
+import os
 from argparse import ArgumentParser
 import onnx
-from typing import List
+from typing import Optional, List
 
 class Color:
     BLACK          = '\033[30m'
@@ -32,20 +33,17 @@ class Color:
 
 def extraction(
     input_onnx_file_path: str,
-    output_onnx_file_path: str,
     input_op_names: List[str],
     output_op_names: List[str],
-):
+    output_onnx_file_path: Optional[str] = '',
+) -> onnx.ModelProto:
 
     """
     Parameters
     ----------
     input_onnx_file_path: str
         Input onnx file path.
 
-    output_onnx_file_path: str
-        Output onnx file path.
-
     input_op_names: List[str]
         List of OP names to specify for the input layer of the model.\n\
         Specify the name of the OP, separated by commas.\n\
@@ -55,15 +53,37 @@ def extraction(
         List of OP names to specify for the output layer of the model.\n\
         Specify the name of the OP, separated by commas.\n\
         e.g. ['ddd','eee','fff']
+
+    output_onnx_file_path: Optional[str]
+        Output onnx file path.\n\
+        If not specified, .onnx is not output.\n\
+        Default: ''
+
+    Returns
+    -------
+    extracted_graph: onnx.ModelProto
+        Extracted onnx ModelProto
     """
 
+    tmp_onnx_file = ''
+    if not output_onnx_file_path:
+        tmp_onnx_file = 'extracted.onnx'
+    else:
+        tmp_onnx_file = output_onnx_file_path
+
     onnx.utils.extract_model(
         input_onnx_file_path,
-        output_onnx_file_path,
+        tmp_onnx_file,
         input_op_names,
         output_op_names
     )
 
+    extracted_graph = onnx.load(tmp_onnx_file)
+    if not output_onnx_file_path:
+        os.remove(tmp_onnx_file)
+
+    return extracted_graph
+
 
 def main():
     parser = ArgumentParser()
@@ -73,12 +93,6 @@ def main():
         required=True,
         help='Input onnx file path.'
     )
-    parser.add_argument(
-        '--output_onnx_file_path',
-        type=str,
-        required=True,
-        help='Output onnx file path.'
-    )
     parser.add_argument(
         '--input_op_names',
         type=str,
@@ -97,18 +111,25 @@ def main():
             Specify the name of the OP, separated by commas. \
             e.g. --output_op_names ddd,eee,fff"
     )
+    parser.add_argument(
+        '--output_onnx_file_path',
+        type=str,
+        default='',
+        help='Output onnx file path. If not specified, extracted.onnx is output.'
+    )
     args = parser.parse_args()
 
     input_op_names = args.input_op_names.strip(' ,').replace(' ','').split(',')
     output_op_names = args.output_op_names.strip(' ,').replace(' ','').split(',')
 
     # Model extraction
-    extraction(
+    extracted_graph = extraction(
         input_onnx_file_path=args.input_onnx_file_path,
-        output_onnx_file_path=args.output_onnx_file_path,
         input_op_names=input_op_names,
         output_op_names=output_op_names,
+        output_onnx_file_path=args.output_onnx_file_path,
     )
 
+
 if __name__ == '__main__':
     main()