fix: update importing onnx module with conditional checks (#105)

* fix: update importing onnx module with conditional checks

* fix: update the version also in init file

CHANGELOG.md

@@ -2,6 +2,12 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+### [1.5.8]
+
+#### Fix
+
+- Fix ONNX import call for the utilities when ONNX is not installed.
+
 ### [1.5.7]
 
 #### Added

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "quadra"
-version = "1.5.7"
+version = "1.5.8"
 description = "Deep Learning experiment orchestration library"
 authors = [
   "Federico Belotti <[email protected]>",

quadra/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.5.7"
+__version__ = "1.5.8"
 
 
 def get_version():

quadra/utils/utils.py

@@ -14,7 +14,6 @@
 import dotenv
 import mlflow
 import numpy as np
-import onnx
 import pytorch_lightning as pl
 import rich.syntax
 import rich.tree
@@ -31,6 +30,14 @@
 from quadra.callbacks.mlflow import get_mlflow_logger
 from quadra.utils.mlflow import infer_signature_model
 
+try:
+    import onnx  # noqa
+
+    ONNX_AVAILABLE = True
+except ImportError:
+    ONNX_AVAILABLE = False
+
+
 IMAGE_EXTENSIONS: List[str] = [".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".pbm", ".pgm", ".ppm", ".pxm", ".pnm"]
 
 
@@ -311,7 +318,7 @@ def finish(
                                     artifact_path=model_path,
                                     signature=signature,
                                 )
-                        elif model_type in ["onnx", "simplified_onnx"]:
+                        elif model_type in ["onnx", "simplified_onnx"] and ONNX_AVAILABLE:
                             signature = infer_signature_model(model, inputs)
                             with mlflow.start_run(run_id=mlflow_logger.run_id) as _:
                                 if model.model_path is None: