fix ImageReader to work either with image path or np.ndarray

Author: pythonlessons

CHANGELOG.md

@@ -1,3 +1,9 @@
+## [1.0.4] - 2022-03-22
+### Changed
+- Fix `ImageReader` to work either with image path or `np.ndarray`
+- Added `metadata` support to `callbacks/tf2onnx` when converting to onnx format
+
+
 ## [1.0.3] - 2022-03-20
 ### Changed
 - Changed `mltu.augmentors` to work only with `Image` objects

mltu/__init__.py

@@ -1,3 +1,3 @@
-__version__ = "1.0.3"
+__version__ = "1.0.4"
 
 from .annotations.image import Image

mltu/annotations/image.py

@@ -17,7 +17,7 @@ def __init__(
         image: typing.Union[str, np.ndarray],
         method: int = cv2.IMREAD_COLOR,
         path: str = "",
-        color: str = ""
+        color: str = "BGR"
         ) -> None:
 
         if isinstance(image, str):

mltu/preprocessors.py

@@ -20,10 +20,25 @@ def __init__(self, method: int = cv2.IMREAD_COLOR, log_level: int = logging.INFO
         self.logger = logging.getLogger(self.__class__.__name__)
         self.logger.setLevel(log_level)
 
-    def __call__(self, image_path: str, label: typing.Any) -> typing.Tuple[Image, typing.Any]:
-        # check whether image_path exists
-        if not os.path.exists(image_path):
-            raise FileNotFoundError(f"Image {image_path} not found.")
+    def __call__(self, image_path: typing.Union[str, np.ndarray], label: typing.Any) -> typing.Tuple[Image, typing.Any]:
+        """ Read image with cv2 from path and return image and label
+        
+        Args:
+            image_path (typing.Union[str, np.ndarray]): Path to image or numpy array
+            label (Any): Label of image
+
+        Returns:
+            Image: Image object
+            Any: Label of image
+        """
+        if isinstance(image_path, str):
+            # check whether image_path exists
+            if not os.path.exists(image_path):
+                raise FileNotFoundError(f"Image {image_path} not found.")
+        elif isinstance(image_path, np.ndarray):
+            pass
+        else:
+            raise TypeError(f"Image {image_path} is not a string or numpy array.")
 
         image = Image(image = image_path, method = self._method)
         if image.image is None:

mltu/tensorflow/callbacks.py

@@ -1,26 +1,43 @@
 import os
+import tf2onnx
+import onnx
 from keras.callbacks import Callback
 
 import logging
 
 class Model2onnx(Callback):
-    """Converts the model to onnx format after training is finished.
-    
-    Args:
-        saved_model_path (str): Path to the saved .h5 model.
-    """
-    try:
-        import tf2onnx
-    except ImportError:
-        raise ImportError("tf2onnx not installed, skipping model export to onnx")
-
-    def __init__(self, saved_model_path: str) -> None:
+    """ Converts the model to onnx format after training is finished. """
+    def __init__(
+        self, 
+        saved_model_path: str, 
+        metadata: dict=None
+        ) -> None:
+        """ Converts the model to onnx format after training is finished.
+        Args:
+            saved_model_path (str): Path to the saved .h5 model.
+            metadata (dict, optional): Dictionary containing metadata to be added to the onnx model. Defaults to None.
+        """
         super().__init__()
         self.saved_model_path = saved_model_path
+        self.metadata = metadata
 
     def on_train_end(self, logs=None):
         self.model.load_weights(self.saved_model_path)
-        self.tf2onnx.convert.from_keras(self.model, output_path=self.saved_model_path.replace(".h5", ".onnx"), )
+        self.onnx_model_path = self.saved_model_path.replace(".h5", ".onnx")
+        self.tf2onnx.convert.from_keras(self.model, output_path=self.onnx_model_path)
+
+        if self.metadata and isinstance(self.metadata, dict):
+            # Load the ONNX model
+            onnx_model = onnx.load(self.onnx_model_path)
+
+            # Add the metadata dictionary to the model's metadata_props attribute
+            for key, value in self.metadata.items():
+                meta = onnx_model.metadata_props.add()
+                meta.key = key
+                meta.value = value
+
+            # Save the modified ONNX model
+            onnx.save(onnx_model, self.onnx_model_path)
 
 class TrainLogger(Callback):
     """Logs training metrics to a file.