From dd4373b9545a1cd295efa9df4212c05447935613 Mon Sep 17 00:00:00 2001 From: Ashay Rane Date: Wed, 11 Oct 2023 07:28:58 -0500 Subject: [PATCH] Allow pathlib.Path values as input to Model construction (#54) As title --- onnx_tool/model.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/onnx_tool/model.py b/onnx_tool/model.py index 6fef606..7a1bec9 100644 --- a/onnx_tool/model.py +++ b/onnx_tool/model.py @@ -1,4 +1,5 @@ import os +import pathlib import onnx @@ -6,10 +7,13 @@ 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)