Skip to content

Commit

Permalink
Allow pathlib.Path values as input to Model construction (ThanatosShi…
Browse files Browse the repository at this point in the history
…nji#54)

As title
ashay authored Oct 11, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
1 parent 9c71ba5 commit dd4373b
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions onnx_tool/model.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,19 @@
import os
import pathlib

import onnx

from .graph import Graph


class Model:
def __init__(self, m: [str, onnx.ModelProto], verbose=False, constant_folding: bool = True,
def __init__(self, m: [str, onnx.ModelProto, pathlib.Path], verbose=False, constant_folding: bool = True,
noderename: bool = False):
self.modelname = ''
if isinstance(m, str):
if isinstance(m, pathlib.Path):
self.modelname = m.name.stem
m = onnx.load_model(m)
elif isinstance(m, str):
self.modelname = os.path.basename(m)
self.modelname = os.path.splitext(self.modelname)[0]
m = onnx.load_model(m)

0 comments on commit dd4373b

Please sign in to comment.