issue with using albumentations with batched mapped tf.data.Dataset object

🐛 Bug

The code was taken from the examples section in albumentations documentation. I modified the code in order to vectorize the map on tf.data.Dataset object. In order to do that I created a resized_ds that contains tuples of (image,label) where each image is of size (120,120). When I run the

ds_alb = resized_ds.batch(32).map(partial(process_data, img_size=120),
                  num_parallel_calls=AUTOTUNE,deterministic=False).prefetch(AUTOTUNE)

(which is the last chunk of code in the notebook I linked here) throws an unknown error sometimes while it works the other times without changing anything in the code

To Reproduce

Steps to reproduce the behavior:

Attaching the link to the colab notebook

Error log:

/usr/local/lib/python3.6/dist-packages/albumentations/augmentations/functional.py:789: UserWarning: Image compression augmentation is most effective with uint8 inputs, float32 is used as input.
  UserWarning,
---------------------------------------------------------------------------
UnknownError                              Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/context.py in execution_mode(mode)
   2101       ctx.executor = executor_new
-> 2102       yield
   2103     finally:

10 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/iterator_ops.py in _next_internal(self)
    757             output_types=self._flat_output_types,
--> 758             output_shapes=self._flat_output_shapes)
    759 

/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gen_dataset_ops.py in iterator_get_next(iterator, output_types, output_shapes, name)
   2609     except _core._NotOkStatusException as e:
-> 2610       _ops.raise_from_not_ok_status(e, name)
   2611     except _core._FallbackException:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)
   6842   # pylint: disable=protected-access
-> 6843   six.raise_from(core._status_to_exception(e.code, message), None)
   6844   # pylint: enable=protected-access

/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

UnknownError: error: OpenCV(4.1.2) /io/opencv/modules/imgcodecs/src/grfmt_base.cpp:145: error: (-10:Unknown error code -10) Raw image encoder error: Maximum supported image dimension is 65500 pixels in function 'throwOnEror'

Traceback (most recent call last):

  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/script_ops.py", line 244, in __call__
    ret = func(*args)

  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/autograph/impl/api.py", line 302, in wrapper
    return func(*args, **kwargs)

  File "<ipython-input-24-d518d3df69c5>", line 3, in aug_fn
    aug_data = transforms(**data)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/core/composition.py", line 176, in __call__
    data = t(force_apply=force_apply, **data)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/core/transforms_interface.py", line 87, in __call__
    return self.apply_with_params(params, **kwargs)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/core/transforms_interface.py", line 100, in apply_with_params
    res[key] = target_function(arg, **dict(params, **target_dependencies))

  File "/usr/local/lib/python3.6/dist-packages/albumentations/augmentations/transforms.py", line 1630, in apply
    return F.image_compression(image, quality, image_type)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/augmentations/functional.py", line 54, in wrapped_function
    result = func(img, *args, **kwargs)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/augmentations/functional.py", line 796, in image_compression
    _, encoded_img = cv2.imencode(image_type, img, (int(quality_flag), quality))

cv2.error: OpenCV(4.1.2) /io/opencv/modules/imgcodecs/src/grfmt_base.cpp:145: error: (-10:Unknown error code -10) Raw image encoder error: Maximum supported image dimension is 65500 pixels in function 'throwOnEror'



	 [[{{node PyFunc}}]] [Op:IteratorGetNext]

During handling of the above exception, another exception occurred:

UnknownError                              Traceback (most recent call last)
<ipython-input-46-ef3ac52e7aa5> in <module>()
      3                   num_parallel_calls=AUTOTUNE,deterministic=False).prefetch(AUTOTUNE)
      4 it = iter(ds_alb)
----> 5 batch = next(it)
      6 print(type(batch))
      7 images,labels = batch

/usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/iterator_ops.py in __next__(self)
    734 
    735   def __next__(self):  # For Python 3 compatibility
--> 736     return self.next()
    737 
    738   def _next_internal(self):

/usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/iterator_ops.py in next(self)
    770   def next(self):
    771     try:
--> 772       return self._next_internal()
    773     except errors.OutOfRangeError:
    774       raise StopIteration

/usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/iterator_ops.py in _next_internal(self)
    762         return self._element_spec._from_compatible_tensor_list(ret)  # pylint: disable=protected-access
    763       except AttributeError:
--> 764         return structure.from_compatible_tensor_list(self._element_spec, ret)
    765 
    766   @property

/usr/lib/python3.6/contextlib.py in __exit__(self, type, value, traceback)
     97                 value = type()
     98             try:
---> 99                 self.gen.throw(type, value, traceback)
    100             except StopIteration as exc:
    101                 # Suppress StopIteration *unless* it's the same exception that

/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/context.py in execution_mode(mode)
   2103     finally:
   2104       ctx.executor = executor_old
-> 2105       executor_new.wait()
   2106 
   2107 

/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/executor.py in wait(self)
     65   def wait(self):
     66     """Waits for ops dispatched in this executor to finish."""
---> 67     pywrap_tfe.TFE_ExecutorWaitForAllPendingNodes(self._handle)
     68 
     69   def clear_error(self):

UnknownError: error: OpenCV(4.1.2) /io/opencv/modules/imgcodecs/src/grfmt_base.cpp:145: error: (-10:Unknown error code -10) Raw image encoder error: Maximum supported image dimension is 65500 pixels in function 'throwOnEror'

Traceback (most recent call last):

  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/script_ops.py", line 244, in __call__
    ret = func(*args)

  File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/autograph/impl/api.py", line 302, in wrapper
    return func(*args, **kwargs)

  File "<ipython-input-24-d518d3df69c5>", line 3, in aug_fn
    aug_data = transforms(**data)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/core/composition.py", line 176, in __call__
    data = t(force_apply=force_apply, **data)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/core/transforms_interface.py", line 87, in __call__
    return self.apply_with_params(params, **kwargs)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/core/transforms_interface.py", line 100, in apply_with_params
    res[key] = target_function(arg, **dict(params, **target_dependencies))

  File "/usr/local/lib/python3.6/dist-packages/albumentations/augmentations/transforms.py", line 1630, in apply
    return F.image_compression(image, quality, image_type)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/augmentations/functional.py", line 54, in wrapped_function
    result = func(img, *args, **kwargs)

  File "/usr/local/lib/python3.6/dist-packages/albumentations/augmentations/functional.py", line 796, in image_compression
    _, encoded_img = cv2.imencode(image_type, img, (int(quality_flag), quality))

cv2.error: OpenCV(4.1.2) /io/opencv/modules/imgcodecs/src/grfmt_base.cpp:145: error: (-10:Unknown error code -10) Raw image encoder error: Maximum supported image dimension is 65500 pixels in function 'throwOnEror'



	 [[{{node PyFunc}}]]

Expected behavior

It was expected that the mapping would apply and could be vectorised without any issues as in case of tf.image based transformations

Environment

• Albumentations version (e.g., 0.1.8): 0.4.6
• Python version (e.g., 3.7): 3.6.9, 3.7.0, 3.8.5
• OS (e.g., Linux): Windows 10
• How you installed albumentations (conda, pip, source): conda
• Any other relevant information:

Additional context

[Dipet]

@thingumajig solution resolves the problem. Also you can rewrite your function that calls albumentations pipeline like this:

def aug_fn(image, img_size):
    data = {"image":image}
    if image.ndim == 4:
        aug_data = np.stack([transforms(image=i) for i in image], axis=0)
    else:
        aug_data = transforms(**data)
    aug_img = aug_data["image"]
    aug_img = tf.cast(aug_img/255.0, tf.float32)
    aug_img = tf.image.resize(aug_img, size=[img_size, img_size])
    
    return aug_img

[AkshayRoyal]

Thanks for the advice but I request the developers to add this feature so that we can directly apply mapping to the batched dataset.