From 43256977faa5bd113906cf32162ee9adb25a1d89 Mon Sep 17 00:00:00 2001
From: "q.yao" <streetyao@live.com>
Date: Wed, 15 Dec 2021 18:43:09 +0800
Subject: [PATCH] Add onnx2ncnn tool (#276)

* Add onnx2ncnn tool

* fix docs

* fix lint
---
 docs/useful_tools.md | 21 +++++++++++++++++++++
 tools/onnx2ncnn.py   | 45 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 tools/onnx2ncnn.py

diff --git a/docs/useful_tools.md b/docs/useful_tools.md
index 605cbea31..cf99502f8 100644
--- a/docs/useful_tools.md
+++ b/docs/useful_tools.md
@@ -84,3 +84,24 @@ python tools/onnx2tensorrt.py \
 - `--device-id` : The device index, default to `0`.
 - `--calib-file` : The calibration data used to calibrate engine to int8.
 - `--log-level` : To set log level which in `'CRITICAL', 'FATAL', 'ERROR', 'WARN', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'`. If not specified, it will be set to `INFO`.
+
+## onnx2ncnn
+
+This tool helps to convert an `ONNX` model to an `ncnn` model.
+
+### Usage
+
+```bash
+python tools/onnx2ncnn.py \
+    ${ONNX_PATH} \
+    ${NCNN_PARAM} \
+    ${NCNN_BIN} \
+    --log-level INFO
+```
+
+### Description of all arguments
+
+- `onnx_path` : The path of the `ONNX` model to convert from.
+- `output_param` : The converted `ncnn` param path.
+- `output_bin` : The converted `ncnn` bin path.
+- `--log-level` : To set log level which in `'CRITICAL', 'FATAL', 'ERROR', 'WARN', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'`. If not specified, it will be set to `INFO`.
diff --git a/tools/onnx2ncnn.py b/tools/onnx2ncnn.py
new file mode 100644
index 000000000..6f54fcc31
--- /dev/null
+++ b/tools/onnx2ncnn.py
@@ -0,0 +1,45 @@
+import argparse
+import logging
+
+from mmdeploy.apis.ncnn import onnx2ncnn
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Convert ONNX to ncnn.')
+    parser.add_argument('onnx_path', help='ONNX model path')
+    parser.add_argument('output_param', help='output ncnn param path')
+    parser.add_argument('output_bin', help='output bin path')
+    parser.add_argument(
+        '--log-level',
+        help='set log level',
+        default='INFO',
+        choices=list(logging._nameToLevel.keys()))
+    args = parser.parse_args()
+
+    return args
+
+
+def main():
+    args = parse_args()
+    logging.basicConfig(
+        format='%(asctime)s,%(name)s %(levelname)-8s'
+        ' [%(filename)s:%(lineno)d] %(message)s',
+        datefmt='%Y-%m-%d:%H:%M:%S')
+    logger = logging.getLogger()
+    logger.setLevel(args.log_level)
+
+    onnx_path = args.onnx_path
+    output_param = args.output_param
+    output_bin = args.output_bin
+
+    logging.info(f'onnx2ncnn: \n\tonnx_path: {onnx_path} ')
+    try:
+        onnx2ncnn(onnx_path, output_param, output_bin)
+        logging.info('onnx2ncnn success.')
+    except Exception as e:
+        logging.error(e)
+        logging.error('onnx2ncnn failed.')
+
+
+if __name__ == '__main__':
+    main()