two new augmentors and some minor bug fixes

Author: pythonlessons

CHANGELOG.md

@@ -1,3 +1,12 @@
+## [1.0.12] - 2022-06-08
+### Changed
+- Moved `onnx` and `tf2onnx` import inside `mltu.tensorflow.callbacks.Model2onnx` to avoid import errors when not using this callback
+- Removed `onnx` and `tf2onnx` install requirements from global requirements
+
+### Added
+- Added `RandomMirror` and `RandomFlip` augmentors into `mltu.augmentors`
+- Added `u2net` segmentation model into `mltu.tensorflow.models`
+
 ## [1.0.11] - 2022-06-07
 ### Changed
 - Downgrade `tf2onnx` and `onnx` versions, they don't work with newest TensorFlow version

mltu/augmentors.py

@@ -12,6 +12,8 @@
 - RandomSharpen
 - RandomGaussianBlur
 - RandomSaltAndPepper
+- RandomMirror
+- RandomFlip
 """
 
 
@@ -177,6 +179,11 @@ def __call__(self, image: Image, annotation: typing.Any) -> typing.Tuple[Image,
 
         # perform the actual rotation and return the image
         img = cv2.warpAffine(image.numpy(), M, (nW, nH), borderValue=borderValue)
+
+        if isinstance(annotation, Image):
+            annotation_warp = cv2.warpAffine(annotation.numpy(), M, (nW, nH), borderValue=(0, 0, 0))
+            annotation.update(annotation_warp)
+
         image.update(img)
 
         return image, annotation
@@ -377,4 +384,72 @@ def __call__(self, image: Image, annotation: typing.Any) -> typing.Tuple[Image,
 
         image.update(img)
 
+        return image, annotation
+    
+
+class RandomMirror(Augmentor):
+    """ Randomly mirror image"""
+    def __init__(
+        self, 
+        random_chance: float = 0.5,
+        log_level: int = logging.INFO,
+        ) -> None:
+        """ Randomly mirror image
+        
+        Args:
+            random_chance (float): Float between 0.0 and 1.0 setting bounds for random probability. Defaults to 0.5.
+            log_level (int): Log level for the augmentor. Defaults to logging.INFO.
+        """
+        super(RandomMirror, self).__init__(random_chance, log_level)
+
+    @randomness_decorator
+    def __call__(self, image: Image, annotation: typing.Any) -> typing.Tuple[Image, typing.Any]:
+        """ Randomly mirror an image
+
+        Args:
+            image (Image): Image to be mirrored
+            annotation (typing.Any): Annotation to be mirrored
+
+        Returns:
+            image (Image): Mirrored image
+            annotation (typing.Any): Mirrored annotation if necessary
+        """
+        image = image.flip(0)
+        if isinstance(annotation, Image):
+            annotation = annotation.flip(0)
+
+        return image, annotation
+    
+
+class RandomFlip(Augmentor):
+    """ Randomly flip image"""
+    def __init__(
+        self, 
+        random_chance: float = 0.5,
+        log_level: int = logging.INFO,
+        ) -> None:
+        """ Randomly mirror image
+        
+        Args:
+            random_chance (float): Float between 0.0 and 1.0 setting bounds for random probability. Defaults to 0.5.
+            log_level (int): Log level for the augmentor. Defaults to logging.INFO.
+        """
+        super(RandomFlip, self).__init__(random_chance, log_level)
+
+    @randomness_decorator
+    def __call__(self, image: Image, annotation: typing.Any) -> typing.Tuple[Image, typing.Any]:
+        """ Randomly mirror an image
+
+        Args:
+            image (Image): Image to be flipped
+            annotation (typing.Any): Annotation to be flipped
+
+        Returns:
+            image (Image): Flipped image
+            annotation (typing.Any): Flipped annotation if necessary
+        """
+        image = image.flip(1)
+        if isinstance(annotation, Image):
+            annotation = annotation.flip(1)
+
         return image, annotation

mltu/preprocessors.py

@@ -13,7 +13,7 @@
 
 class ImageReader:
     """Read image from path and return image and label"""
-    def __init__(self, image_class, log_level: int = logging.INFO, ) -> None:
+    def __init__(self, image_class: Image, log_level: int = logging.INFO, ) -> None:
         self.logger = logging.getLogger(self.__class__.__name__)
         self.logger.setLevel(log_level)
         self._image_class = image_class

mltu/tensorflow/callbacks.py

@@ -4,6 +4,11 @@
 import logging
 
 class Model2onnx(Callback):
+
+    # import onnx and use as global
+    import onnx
+    import tf2onnx
+
     """ Converts the model to onnx format after training is finished. """
     def __init__(
         self, 
@@ -22,15 +27,15 @@ def __init__(
     def on_train_end(self, logs=None):
         """ Converts the model to onnx format after training is finished. """
         try:
-            import onnx
-            import tf2onnx
+            # import onnx
+            # import tf2onnx
             self.model.load_weights(self.saved_model_path)
             self.onnx_model_path = self.saved_model_path.replace(".h5", ".onnx")
-            tf2onnx.convert.from_keras(self.model, output_path=self.onnx_model_path)
+            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)
+                onnx_model = self.onnx.load(self.onnx_model_path)
 
                 # Add the metadata dictionary to the model's metadata_props attribute
                 for key, value in self.metadata.items():
@@ -39,7 +44,7 @@ def on_train_end(self, logs=None):
                     meta.value = value
 
                 # Save the modified ONNX model
-                onnx.save(onnx_model, self.onnx_model_path)
+                self.onnx.save(onnx_model, self.onnx_model_path)
 
         except Exception as e:
             print(e)