Updated data example notebook and added better backend switching

Author: amogh7joshi

agml/data/managers/transform_helpers.py

@@ -145,3 +145,17 @@ def apply(self, mask):
         return out[..., 1:].astype(np.int32)
 
 
+class ToTensorSeg(TransformApplierBase):
+    def apply(self, image, mask):
+        # This is essentially the same as `torchvision.transforms.ToTensor`,
+        # except it doesn't scale the values of the segmentation mask from
+        # 0-255 -> 0-1, as this would break the segmentaiotn pipeline.
+        if image.ndim == 2:
+            image = image[:, :, None]
+        image = torch.from_numpy(image).permute(2, 0, 1).contiguous()
+        if isinstance(image, torch.ByteTensor):
+            image = image.to(dtype = torch.float32).div(255)
+        mask = torch.from_numpy(mask)
+        return mask
+
+

agml/data/managers/transforms.py

@@ -30,7 +30,8 @@
     ScaleTransform,
     NormalizationTransform,
     OneHotLabelTransform,
-    MaskToChannelBasisTransform
+    MaskToChannelBasisTransform,
+    ToTensorSeg
 )
 from agml.utils.logging import log
 
@@ -67,7 +68,8 @@ class TransformManager(AgMLSerializable):
     there is no `dual_transform` for image classification.
     """
     serializable = frozenset(
-        ('task', 'transforms', 'time_inserted_transforms'))
+        ('task', 'transforms', 'time_inserted_transforms',
+         'contains_tf_transforms', 'contains_torch_transforms'))
 
     def __init__(self, task):
         self._task = task
@@ -87,6 +89,12 @@ def __init__(self, task):
         # in this list and applies them as required.
         self._time_inserted_transforms = []
 
+        # Check if the loader already contains TensorFlow/PyTorch transforms.
+        # This is to track whether there are issues when a new transform is added
+        # or if there already exists one, so we whether to apply transforms.
+        self._contains_tf_transforms = False
+        self._contains_torch_transforms = False
+
     def get_transform_states(self):
         """Returns a copy of the existing transforms."""
         transform_dict = {}
@@ -319,13 +327,34 @@ def _build_normalization_transform(self, transform):
                 self._time_inserted_transforms.pop(norm_transform_index_time)
         return None
 
+    def _transform_update_and_check(self, tfm_type):
+        """Checks whether a TensorFlow/PyTorch transform has been added."""
+        # Check for transform conflicts.
+        if tfm_type == 'torch' and self._contains_tf_transforms:
+            raise TypeError("Received a PyTorch-type transform, yet the loader "
+                            "already contains TensorFlow/Keras transforms. This "
+                            "will cause an error, please only pass one format. If "
+                            "you want to remove a transform, pass a value of `None` "
+                            "to reset all of the transforms for a certain type.")
+        if tfm_type == 'tf' and self._contains_torch_transforms:
+            raise TypeError("Received a TensorFlow/Keras-type transform, yet the "
+                            "loader already contains PyTorch transforms. This "
+                            "will cause an error, please only pass one format. If "
+                            "you want to remove a transform, pass a value of `None` "
+                            "to reset all of the transforms for a certain type.")
+
+        # Update the transform type.
+        if tfm_type == 'torch':
+            self._contains_torch_transforms = True
+        if tfm_type == 'tf':
+            self._contains_tf_transforms = True
+
     # The following methods implement different checks which validate
     # as well as process input transformations, and manage the backend.
     # The transforms here will be also checked to match a specific
     # backend. Alternatively, the backend will dynamically be switched.
 
-    @staticmethod
-    def _construct_single_image_transform(transform):
+    def _construct_single_image_transform(self, transform):
         """Validates a transform which is applied to a single image.
 
         This is used for image classification transforms, which only
@@ -361,6 +390,7 @@ def _construct_single_image_transform(transform):
                 if user_changed_backend():
                     raise StrictBackendError(change = 'tf', obj = transform)
                 set_backend('torch')
+            self._transform_update_and_check('torch')
             return transform
 
         # A `tf.keras.Sequential` preprocessing model or an individual
@@ -370,15 +400,15 @@ def _construct_single_image_transform(transform):
                 if user_changed_backend():
                     raise StrictBackendError(change = 'torch', obj = transform)
                 set_backend('tf')
+            self._transform_update_and_check('tf')
             return transform
 
         # Otherwise, it may be a transform from a (lesser-known) third-party
         # library, in which case we just return it as a callable. Transforms
         # which are used in a more complex manner should be passed as decorators.
         return transform
 
-    @staticmethod
-    def _construct_image_and_mask_transform(transform):
+    def _construct_image_and_mask_transform(self, transform):
         """Validates a transform for an image and annotation mask.
 
         This is used for a semantic segmentation transform. Such
@@ -426,18 +456,41 @@ def _construct_image_and_mask_transform(transform):
                     if user_changed_backend():
                         raise StrictBackendError(
                             change = 'tf', obj = transform)
-                    set_backend('tf')
+                    set_backend('torch')
+
+                # Update `torchvision.transforms.ToTensor` to a custom
+                # updated class as this will modify the mask incorrectly.
+                import torchvision
+                if isinstance(transform, torchvision.transforms.ToTensor):
+                    transform = ToTensorSeg(None)
+                    log("Updated `ToTensor` transform in the provided pipeline "
+                        f"{transform} to an updated transform which does not "
+                        f"modify the mask. If you want to change this behaviour, "
+                        f"please raise an error with the AgML team.")
+                elif isinstance(transform, torchvision.transforms.Compose):
+                    tfm_list = transform.transforms.copy()
+                    for i, compose_tfm in enumerate(transform.transforms):
+                        if isinstance(compose_tfm, torchvision.transforms.ToTensor):
+                            tfm_list[i] = ToTensorSeg(None)
+                    transform = torchvision.transforms.Compose(tfm_list)
+                    log("Updated `ToTensor` transform in the provided pipeline "
+                        f"{transform} to an updated transform which does not "
+                        f"modify the mask. If you want to change this behaviour, "
+                        f"please raise an error with the AgML team.")
+
+                self._transform_update_and_check('torch')
             elif 'keras.layers' in transform.__module__:
                 if get_backend() != 'tf':
                     if user_changed_backend():
                         raise StrictBackendError(
                             change = 'torch', obj = transform)
-                    set_backend('torch')
+                    set_backend('tf')
                 log('Got a Keras transformation for a dual image and '
                     'mask transform. If you are passing preprocessing '
                     'layers to this method, then use `agml.data.experimental'
                     '.generate_keras_segmentation_dual_transform` in order '
                     'for the random state to be applied properly.', 'warning')
+                self._transform_update_and_check('tf')
             return SameStateImageMaskTransform(transform)
 
         # Another type of transform, most likely some form of transform

agml/viz/__init__.py

@@ -24,7 +24,7 @@
     convert_mask_to_colored_image,
     annotate_semantic_segmentation,
     show_image_and_mask,
-    show_image_with_overlaid_mask,
+    show_image_and_overlaid_mask,
     show_semantic_segmentation_truth_and_prediction,
 )
 from .boxes import (

agml/viz/general.py

@@ -16,7 +16,7 @@
 
 from agml.backend.tftorch import is_array_like
 from agml.viz.boxes import show_image_and_boxes
-from agml.viz.masks import show_image_with_overlaid_mask
+from agml.viz.masks import show_image_and_overlaid_mask
 from agml.viz.labels import show_images_and_labels
 from agml.viz.tools import format_image, _inference_best_shape, convert_figure_to_image
 from agml.viz.display import display_image
@@ -52,7 +52,7 @@ def show_sample(loader, image_only = False, **kwargs):
         return show_image_and_boxes(
             sample, info = loader.info, no_show = kwargs.get('no_show', False))
     elif loader.task == 'semantic_segmentation':
-        return show_image_with_overlaid_mask(
+        return show_image_and_overlaid_mask(
             sample, no_show = kwargs.get('no_show', False))
     elif loader.task == 'image_classification':
         return show_images_and_labels(

agml/viz/masks.py

@@ -28,9 +28,9 @@ def binary_to_channel_by_channel(mask, num_classes = None):
     if mask.ndim == 3 and not np.all(mask[:, :, 0] == mask[:, :, 1]):
         return mask
 
-    # For binary classification tasks, return the mask as-is.
+    # For binary classification tasks, return the mask with an additional channel.
     if len(np.unique(mask)) == 2 and mask.max() == 1:
-        return mask
+        return np.expand_dims(mask, -1)
 
     # Otherwise, convert the mask to channel-by-channel format.
     input_shape = mask.shape
@@ -137,11 +137,11 @@ def annotate_semantic_segmentation(image,
     return image
 
 
-def show_image_with_overlaid_mask(image,
-                                  mask = None,
-                                  alpha = 0.3,
-                                  border = True,
-                                  **kwargs):
+def show_image_and_overlaid_mask(image,
+                                 mask = None,
+                                 alpha = 0.3,
+                                 border = True,
+                                 **kwargs):
     """Displays an image with an annotated segmentation mask.
 
     This method overlays a segmentation mask over an image. It uses contours