Image decoding often fails when  transforming Image datasets
Opened by RafayAK (closed)

## Describe the bug
When transforming/modifying images in an image dataset using the `map` function the PIL images often fail to decode in time for the image transforms, causing errors.

Using a debugger it is easy to see what the problem is, the Image decode invocation does not take place and the resulting image passed around is still raw bytes:
```
[{'bytes': b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00 \x00\x00\x00 \x08\x02\x00\x00\x00\xfc\x18\xed\xa3\x00\x00\x08\x02IDATx\x9cEVIs[\xc7\x11\xeemf\xde\x82\x8d\x80\x08\x89"\xb5V\\\xb6\x94(\xe5\x9f\x90\xca5\x7f$\xa7T\xe5\x9f&9\xd9\x8a\\.\xdb\xa4$J\xa4\x00\x02x\xc0{\xb3t\xe7\x00\xca\x99\xd3\\f\xba\xba\xbf\xa5?|\xfa\xf4\xa2\xeb\xba\xedv\xa3f^\xf8\xd5\x0bY\xb6\x10\xb3\xaaDq\xcd\x83\x87\xdf5\xf3gZ\x1a\x04\x0f\xa0fp\xfa\xe0\xd4\x07?\x9dN\xc4\xb1\x99\xfd\xf2\xcb/\x97\x97\x97H\xa2\xaaf\x16\x82\xaf\xeb\xca{\xbf\xd9l.\xdf\x7f\xfa\xcb_\xff&\x88\x08\x00\x80H\xc0\x80@.;\x0f\x8c@#v\xe3\xe5\xfc\xd1\x9f\xee6q\xbf\xdf\xa6\x14\'\x93\xf1\xc3\xe5\xe3\xd1x\x14c\x8c1\xa5\x1c\x9dsM\xd3\xb4\xed\x08\x89SJ)\xa5\xedv\xbb^\xafNO\x97D\x84Hf .... 
```

## Steps to reproduce the bug
```python
from datasets import load_dataset, Dataset
import numpy as np
# seeded NumPy random number generator for reprodducinble results.
rng = np.random.default_rng(seed=0)

test_dataset = load_dataset('cifar100', split="test")

def preprocess_data(dataset):
    """
    Helper function to pre-process HuggingFace Cifar-100 Dataset to remove fine_label and coarse_label columns and
    add is_flipped column
    Args:
        dataset: HuggingFace CIFAR-100 Dataset Object

    Returns:
        new_dataset: A Dataset object with "img" and "is_flipped" columns only

    """
    # remove fine_label and coarse_label columns
    new_dataset = dataset.remove_columns(['fine_label', 'coarse_label'])
    # add the column for is_flipped
    new_dataset = new_dataset.add_column(name="is_flipped", column=np.zeros((len(new_dataset)), dtype=np.uint8))

    return new_dataset


def generate_flipped_data(example, p=0.5):
    """
    A Dataset mapping function that transforms some of the images up-side-down.
    If the probability value (p) is 0.5 approximately half the images will be flipped upside-down
    Args:
        example: An example from the dataset containing a Python dictionary with "img" and "is_flipped" key-value pair
        p: the probability of flipping the image up-side-down, Default 0.5

    Returns:
        example: A Dataset object

    """
    # example['img'] = example['img']
    if rng.random() > p:  # the flip the image and set is_flipped column to 1
        example['img'] = example['img'].transpose(
            1)  # ImageOps.flip(example['img'])  #example['img'].transpose(Image.FLIP_TOP_BOTTOM)
        example['is_flipped'] = 1

    return example

my_test = preprocess_data(test_dataset)
my_test = my_test.map(generate_flipped_data)

```

## Expected results
The dataset should be transformed without problems.

## Actual results
```
/home/rafay/anaconda3/envs/pytorch_new/bin/python /home/rafay/Documents/you_only_live_once/upside_down_detector/create_dataset.py
Reusing dataset cifar100 (/home/rafay/.cache/huggingface/datasets/cifar100/cifar100/1.0.0/f365c8b725c23e8f0f8d725c3641234d9331cd2f62919d1381d1baa5b3ba3142)
Reusing dataset cifar100 (/home/rafay/.cache/huggingface/datasets/cifar100/cifar100/1.0.0/f365c8b725c23e8f0f8d725c3641234d9331cd2f62919d1381d1baa5b3ba3142)
 20%|█▉        | 1999/10000 [00:00<00:01, 5560.44ex/s]
Traceback (most recent call last):
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 2326, in _map_single
    writer.write(example)
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_writer.py", line 441, in write
    self.write_examples_on_file()
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_writer.py", line 399, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_writer.py", line 492, in write_batch
    arrays.append(pa.array(typed_sequence))
  File "pyarrow/array.pxi", line 230, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_writer.py", line 185, in __arrow_array__
    out = pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
  File "pyarrow/array.pxi", line 316, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 143, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 99, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: Could not convert <PIL.Image.Image image mode=RGB size=32x32 at 0x7F56AEE61DE0> with type Image: did not recognize Python value type when inferring an Arrow data type

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/rafay/Documents/you_only_live_once/upside_down_detector/create_dataset.py", line 55, in <module>
    my_test = my_test.map(generate_flipped_data)
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 1953, in map
    return self._map_single(
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 519, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 486, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/fingerprint.py", line 458, in wrapper
    out = func(self, *args, **kwargs)
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 2360, in _map_single
    writer.finalize()
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_writer.py", line 522, in finalize
    self.write_examples_on_file()
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_writer.py", line 399, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_writer.py", line 492, in write_batch
    arrays.append(pa.array(typed_sequence))
  File "pyarrow/array.pxi", line 230, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/home/rafay/anaconda3/envs/pytorch_new/lib/python3.10/site-packages/datasets/arrow_writer.py", line 185, in __arrow_array__
    out = pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
  File "pyarrow/array.pxi", line 316, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 39, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 143, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 99, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: Could not convert <PIL.Image.Image image mode=RGB size=32x32 at 0x7F56AEE61DE0> with type Image: did not recognize Python value type when inferring an Arrow data type

Process finished with exit code 1
```


## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.0.0
- Platform: Linux(Fedora 35)
- Python version: 3.10
- PyArrow version: 7.0.0


Comments: A quick hack I have found is that we can call the image first before running the transforms and it makes sure the image is decoded before being passed on.

For this I just needed to add `example['img'] = example['img']` to the top of my `generate_flipped_data` function, defined above, so that image decode in invoked.

After this minor change this function works:
```python
def generate_flipped_data(example, p=0.5):
    """
    A Dataset mapping functions that transforms some of the image up-side-down.
    If the probability value (p) is 0.5 approximately half the images will be flipped upside-down
    Args:
        example: An example from the dataset containing a Python dictionary with "img" and "is_flipped" key-value pair
        p: probability of flipping the image up-side-down, Default 0.5

    Returns:
        example: A Dataset object

    """
    example['img'] = example['img']  # <<< This is the only change
    if rng.random() > p:  # the flip the image and set is_flipped column to 1
        example['img'] = example['img'].transpose(
            1)  # ImageOps.flip(example['img'])  #example['img'].transpose(Image.FLIP_TOP_BOTTOM)
        example['is_flipped'] = 1

    return example
```

Hi @RafayAK, thanks for reporting.

Current implementation of the Image feature performs the decoding only if the "img" field is accessed by the mapped function.

In your original `generate_flipped_data` function:
- it only accesses the "img" field (and thus performs decoding) if `rng.random() > p`;
- on the other hand, for the cases where `rng.random() <= p`, the "img" field is not accessed and thus no decoding is performed for those examples

By adding the code line `example['img'] = example['img']`, you make sure the "img" field is accessed in all cases, and the decoding is done for all examples.

Also note that there is a little bug in your implementation: `p` is not the probability of flipping, but the probability of not-flipping; the larger is `p`, the smaller is the probability of flipping.

Some refactoring (fixing also `p`):
```python
def generate_flipped_data(example, p=0.5):
    """
    A Dataset mapping functions that transforms some of the image up-side-down.
    If the probability value (p) is 0.5 approximately half the images will be flipped upside-down.

    Args:
        example: An example from the dataset containing a Python dictionary with "img" and "is_flipped" key-value pair
        p: probability of flipping the image up-side-down, Default 0.5

    Returns:
        example: A Dataset object

    """
    do_flip = rng.random() < p  # Note the "<" sign here instead of ">"
    example['img'] = example['img'].transpose(1) if do_flip else example['img']  # Note "img" is always accessed
    example['is_flipped'] = 1 if do_flip else 0
    return example

@albertvillanova Thanks for letting me know this is intended behavior. The docs are severely lacking on this, if I hadn't posted this here I would have never found out how I'm actually supposed to modify images in a Dataset object.

@albertvillanova Secondly if you check the error message it shows that around 1999 images were successfully created, I'm pretty sure some of them were also flipped during the process. Back to my main contention, sometimes the decoding takes place other times it fails. 

I suppose to run `map` on any dataset all the examples should be invoked even if on some of them we end up doing nothing, is that right?

Hi @RafayAK! I've opened a PR with the fix, which adds a fallback to reattempt casting to PyArrow format with a more robust (but more expensive) procedure if the first attempt fails. Feel free to test it by installing `datasets` from the PR branch with the following command:
```
pip install git+https://github.com/huggingface/datasets.git@fix-4124
```

@mariosasko I'll try this right away and report back.

@mariosasko Thanks a lot for looking into this, now the `map` function at least behaves as one would expect a function to behave. 

Looking forward to exploring Hugging Face more and even contributing 😃.

```bash
 $ conda list | grep datasets
datasets                  2.0.1.dev0               pypi_0    pypi

```

```python
def preprocess_data(dataset):
    """
    Helper funtion to pre-process HuggingFace Cifar-100 Dataset to remove fine_label and coarse_label columns and
    add is_flipped column
    Args:
        dataset: HuggingFace CIFAR-100 Dataset Object

    Returns:
        new_dataset: A Dataset object with "img" and "is_flipped" columns only

    """
    # remove fine_label and coarse_label columns
    new_dataset = dataset.remove_columns(['fine_label', 'coarse_label'])
    # add the column for is_flipped
    new_dataset = new_dataset.add_column(name="is_flipped", column=np.zeros((len(new_dataset)), dtype=np.uint8))

    return new_dataset


def generate_flipped_data(example, p=0.5):
    """
    A Dataset mapping functions that transforms some of the image up-side-down.
    If the probability value (p) is 0.5 approximately half the images will be flipped upside-down
    Args:
        example: An example from the dataset containing a Python dictionary with "img" and "is_flipped" key-value pair
        p: probability of flipping the image up-side-down, Default 0.5

    Returns:
        example: A Dataset object

    """
    # example['img'] = example['img']
    if rng.random() > p:  # the flip the image and set is_flipped column to 1
        example['img'] = example['img'].transpose(
            1)  # ImageOps.flip(example['img'])  #example['img'].transpose(Image.FLIP_TOP_BOTTOM)
        example['is_flipped'] = 1

    return example

my_test = preprocess_data(test_dataset)
my_test = my_test.map(generate_flipped_data)
```

The output now show the function was applied successfully:
``` bash
/home/rafay/anaconda3/envs/pytorch_new/bin/python /home/rafay/Documents/you_only_live_once/upside_down_detector/create_dataset.py
Downloading builder script: 5.61kB [00:00, 3.16MB/s]                   
Downloading metadata: 4.21kB [00:00, 2.56MB/s]                   
Reusing dataset cifar100 (/home/rafay/.cache/huggingface/datasets/cifar100/cifar100/1.0.0/f365c8b725c23e8f0f8d725c3641234d9331cd2f62919d1381d1baa5b3ba3142)
Reusing dataset cifar100 (/home/rafay/.cache/huggingface/datasets/cifar100/cifar100/1.0.0/f365c8b725c23e8f0f8d725c3641234d9331cd2f62919d1381d1baa5b3ba3142)
100%|██████████| 10000/10000 [00:01<00:00, 5149.15ex/s]
```