Create onnxvis entry point; fix when main graph does not have a name (#22)

Also fixed a bug when the main graph does not have a name

Author: justinchuby

.gitignore

@@ -158,3 +158,7 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+*.onnx
+*.textproto
+*.pb

README.md

@@ -14,6 +14,12 @@ pip install --upgrade model-explorer-onnx
 
 ```bash
 model-explorer --extensions=model_explorer_onnx
+
+# Or as a shortcut
+onnxvis
+
+# Supply model path
+onnxvis model.onnx
 ```
 
 ## Screenshots
@@ -29,4 +35,3 @@ model-explorer --extensions=model_explorer_onnx
 <img width="1293" alt="image" src="https://github.com/justinchuby/model-explorer-onnx/assets/11205048/fbf2fa05-bd29-4938-93d1-709690d9f9c6">
 
 <img width="1301" alt="image" src="https://github.com/justinchuby/model-explorer-onnx/assets/11205048/a68f7ecd-1fa1-4eac-9e1f-8e9a5bbf9fe3">
-

pyproject.toml

@@ -9,10 +9,10 @@ description = "Adapter for ai-edge-model-explorer to support ONNX models"
 authors = [{ name = "Justin Chu", email = "[email protected]" }]
 readme = "README.md"
 requires-python = ">=3.8"
-license = { file = "LICENSE" }
+license = {text = "MIT License"}
+keywords = ["onnx", "model-explorer", "visualization"]
 classifiers = [
-  "Development Status :: 3 - Alpha",
-  "Environment :: Console",
+  "Development Status :: 4 - Beta",
   "Intended Audience :: Developers",
   "Operating System :: POSIX",
   "Operating System :: MacOS :: MacOS X",
@@ -33,5 +33,8 @@ dependencies = [
   "ml_dtypes",
 ]
 
+[project.scripts]
+onnxvis = "model_explorer_onnx.bin.onnxvis:main"
+
 [project.urls]
 Repository = "https://github.com/justinchuby/model-explorer-onnx"

src/model_explorer_onnx/bin/__init__.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python
+"""A shortcut to run model explorer with ONNX extension."""
+
+import subprocess
+import sys
+
+
+def main():
+    # Run model explorer
+    subprocess.run(
+        [
+            "model-explorer",
+            "--extensions",
+            "model_explorer_onnx",
+            *sys.argv[1:],
+        ]
+    )
+
+
+if __name__ == "__main__":
+    main()

src/model_explorer_onnx/bin/onnxvis.py

@@ -482,10 +482,16 @@ def convert(
         if opset_version is None:
             opset_version = _DEFAULT_OPSET_VERSION
         # TODO: Better support subgraphs in nodes
+        if model.graph.name is None:
+            model.graph.name = "<main>"
+            logger.warning(
+                "Main graph of ONNX file '%s' does not have a name. Set name to '<main>'",
+                model_path,
+            )
         main_graph = create_graph(
             model.graph, all_function_ids, opset_version=opset_version
         )
-        assert main_graph is not None
+        assert main_graph is not None, "Bug: Main graph should not be None"
         graphs.append(main_graph)
 
         for function in model.functions.values():