ImageShowCV2 documentation and OnnxInferenceModel force_cpu flag fix

Author: pythonlessons

CHANGELOG.md

@@ -1,9 +1,11 @@
-## [0.1.5] - 2022-01-03
+## [0.1.5] - 2022-01-10
 
 ### Changed
 - seperated CWERMetric to SER and WER Metrics in mltu.metrics, Character/word rate was calculatted in a wrong way
 - created @setter for augmentors and transformers in DataProvider, to properlly add augmentors and transformers to the pipeline
 - augmentors and transformers must inherit from `mltu.augmentors.base.Augmentor` and `mltu.transformers.base.Transformer` respectively
+- updated ImageShowCV2 transformer documentation
+- fixed OnnxInferenceModel in mltu.inferenceModels to use CPU even if GPU is available with force_cpu=True flag
 
 ### Added:
 - added RandomSharpen to mltu.augmentors, used for simple image augmentation;

mltu/inferenceModel.py

@@ -4,20 +4,20 @@
 import onnxruntime as ort
 
 class OnnxInferenceModel:
-    """ Base class for all inference models that use onnxruntime """
+    """ Base class for all inference models that use onnxruntime 
+
+    Attributes:
+        model_path (str, optional): Path to the model folder. Defaults to "".
+        force_cpu (bool, optional): Force the model to run on CPU or GPU. Defaults to GPU.
+        default_model_name (str, optional): Default model name. Defaults to "model.onnx".
+    """
     def __init__(
         self, 
         model_path: str = "",
         force_cpu: bool = False,
         default_model_name: str = "model.onnx"
         ):
-        """ Initialize the model
 
-        Args:
-            model_path (str, optional): Path to the model. Defaults to "".
-            force_cpu (bool, optional): Force the model to run on CPU or GPU. Defaults to GPU.
-            default_model_name (str, optional): Default model name. Defaults to "model.onnx".
-        """
         self.model_path = model_path
         self.force_cpu = force_cpu
         self.default_model_name = default_model_name
@@ -28,9 +28,14 @@ def __init__(
         if not os.path.exists(self.model_path):
             raise Exception(f"Model path ({self.model_path}) does not exist")
 
-        providers = ['CUDAExecutionProvider'] if ort.get_device() == "GPU" and not force_cpu else ['CPUExecutionProvider']
+        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if ort.get_device() == "GPU" and not force_cpu else ['CPUExecutionProvider']
 
         self.model = ort.InferenceSession(self.model_path, providers=providers)
+
+        # Update providers priority to only CPUExecutionProvider
+        if self.force_cpu:
+            self.model.set_providers(['CPUExecutionProvider'])
+
         self.input_shape = self.model.get_inputs()[0].shape[1:]
         self.input_name = self.model._inputs_meta[0].name
         self.output_name = self.model._outputs_meta[0].name

mltu/transformers.py

@@ -2,6 +2,11 @@
 import typing
 import numpy as np
 
+import logging
+logging.basicConfig(format='%(asctime)s %(levelname)s %(name)s: %(message)s')
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
 class Transformer:
     def __init__(self, *args, **kwargs):
         pass
@@ -89,8 +94,27 @@ def __call__(self, data: np.ndarray, label: np.ndarray):
 
 class ImageShowCV2(Transformer):
     """Show image for visual inspection
+
+    Attributes:
+        verbose (bool): Whether to log label
     """
+    def __init__(self, verbose: bool=True) -> None:
+        self.verbose = verbose
+
     def __call__(self, data: np.ndarray, label: np.ndarray):
+        """ Show image for visual inspection
+
+        Args:
+            data (np.ndarray): Image data
+            label (np.ndarray): Label data
+        
+        Returns:
+            data (np.ndarray): Image data
+            label (np.ndarray): Label data (unchanged)
+        """
+        if self.verbose:
+            logger.info('Label: ', label)
+
         cv2.imshow('image', data)
         cv2.waitKey(0)
         cv2.destroyAllWindows()