convert superpoint onnx

Author: xmba15

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "scripts/superpoint/SuperPointPretrainedNetwork"]
+	path = scripts/superpoint/SuperPointPretrainedNetwork
+	url = https://github.com/magicleap/SuperPointPretrainedNetwork

scripts/superpoint/README.md

@@ -0,0 +1,31 @@
+# convert pre-trained superpoint pytorch weights to onnx format
+
+---
+
+## dependencies
+
+---
+
+- python: 3x
+
+-
+
+```bash
+python3 -m pip install -r requirements.txt
+```
+
+## :running: how to run
+
+---
+
+- update submodule
+
+```bash
+git submodule update --init --recursive
+```
+
+- export onnx weights
+
+```
+python3 convert_to_onnx.py
+```

scripts/superpoint/SuperPointPretrainedNetwork

@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+import os
+
+import torch
+import torch.onnx
+
+from SuperPointPretrainedNetwork.demo_superpoint import SuperPointNet
+
+_CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
+_WEIGHTS_PATH = os.path.join(
+    _CURRENT_DIR, "SuperPointPretrainedNetwork/superpoint_v1.pth"
+)
+
+
+def main():
+    assert os.path.isfile(_WEIGHTS_PATH)
+    model = SuperPointNet()
+    model.load_state_dict(torch.load(_WEIGHTS_PATH))
+    model.eval()
+
+    batch_size = 1
+    height = 16
+    width = 16
+    x = torch.randn(batch_size, 1, height, width)
+
+    torch.onnx.export(
+        model,
+        x,
+        "super_point.onnx",
+        export_params=True,
+        opset_version=10,
+        do_constant_folding=True,
+        input_names=["input"],
+        output_names=["output"],
+        dynamic_axes={
+            "input": {0: "batch_size", 2: "height", 3: "width"},
+            "output": {0: "batch_size"},
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/superpoint/convert_to_onnx.py

@@ -0,0 +1,2 @@
+opencv-python
+torch>=0.4