GitHub issue: Image cast_storage very slow for arrays (e.g. numpy, tensors)
Author: Modexus (CONTRIBUTOR)
State: open | Created: 2024-04-05 13:46:54

Update: see comments below

### Describe the bug

Operations that save an image from a path are very slow.
I believe the reason for this is that the image data (`numpy`) is converted into `pyarrow` format but then back to python using `.pylist()` before being converted to a numpy array again.

`pylist` is already slow but used on a multi-dimensional numpy array such as an image it takes a very long time.

From the trace below we can see that `__arrow_array__` takes a long time.
It is currently also called in `get_inferred_type`, this should be removable #6781 but doesn't change the underyling issue.

The conversion to `pyarrow` and back also leads to the `numpy` array having type `int64` which causes a warning message because the image type excepts `uint8`.
However, originally the `numpy` image array was in `uint8`.

### Steps to reproduce the bug

```python
from PIL import Image
import numpy as np
import datasets
import cProfile

image = Image.fromarray(np.random.randint(0, 255, (2048, 2048, 3), dtype=np.uint8))
image.save("test_image.jpg")

ds = datasets.Dataset.from_dict(
    {"image": ["test_image.jpg"]},
    features=datasets.Features({"image": datasets.Image(decode=True)}),
)

# load as numpy array, e.g. for further processing with map
# same result as map returning numpy arrays
ds.set_format("numpy")

cProfile.run("ds.map(writer_batch_size=1, load_from_cache_file=False)", "restats")
```
```bash
Fri Apr  5 14:56:17 2024    restats

         66817 function calls (64992 primitive calls) in 33.382 seconds

   Ordered by: cumulative time
   List reduced from 1073 to 20 due to restriction <20>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     46/1    0.000    0.000   33.382   33.382 {built-in method builtins.exec}
        1    0.000    0.000   33.382   33.382 <string>:1(<module>)
        1    0.000    0.000   33.382   33.382 arrow_dataset.py:594(wrapper)
        1    0.000    0.000   33.382   33.382 arrow_dataset.py:551(wrapper)
        1    0.000    0.000   33.379   33.379 arrow_dataset.py:2916(map)
        4    0.000    0.000   33.327    8.332 arrow_dataset.py:3277(_map_single)
        1    0.000    0.000   33.311   33.311 arrow_writer.py:465(write)
        2    0.000    0.000   33.311   16.656 arrow_writer.py:423(write_examples_on_file)
        1    0.000    0.000   33.311   33.311 arrow_writer.py:527(write_batch)
        2   14.484    7.242   33.260   16.630 arrow_writer.py:161(__arrow_array__)
        1    0.001    0.001   16.438   16.438 arrow_writer.py:121(get_inferred_type)
        1    0.000    0.000   14.398   14.398 threading.py:637(wait)
        1    0.000    0.000   14.398   14.398 threading.py:323(wait)
        8   14.398    1.800   14.398    1.800 {method 'acquire' of '_thread.lock' objects}
      4/2    0.000    0.000    4.337    2.169 table.py:1800(wrapper)
        2    0.000    0.000    4.337    2.169 table.py:1950(cast_array_to_feature)
        2    0.475    0.238    4.337    2.169 image.py:209(cast_storage)
        9    2.583    0.287    2.583    0.287 {built-in method numpy.array}
        2    0.000    0.000    1.284    0.642 image.py:319(encode_np_array)
        2    0.000    0.000    1.246    0.623 image.py:301(image_to_bytes)
```

### Expected behavior

The `numpy` image data should be passed through as it will be directly consumed by `pillow` to convert it to bytes.

As an example one can replace `list_of_np_array_to_pyarrow_listarray(data)` in `__arrow_array__` with just `out = data` as a test.
We have to change `cast_storage` of the `Image` feature so it handles the passed through data (& if to handle type before)
```python
bytes_array = pa.array(
    [encode_np_array(arr)["bytes"] if arr is not None else None for arr in storage],
    type=pa.binary(),
)
```
Leading to the following:
```bash
Fri Apr  5 15:44:27 2024    restats

         66419 function calls (64595 primitive calls) in 0.937 seconds

   Ordered by: cumulative time
   List reduced from 1023 to 20 due to restriction <20>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     47/1    0.000    0.000    0.935    0.935 {built-in method builtins.exec}
      2/1    0.000    0.000    0.935    0.935 <string>:1(<module>)
      2/1    0.000    0.000    0.934    0.934 arrow_dataset.py:594(wrapper)
      2/1    0.000    0.000    0.934    0.934 arrow_dataset.py:551(wrapper)
      2/1    0.000    0.000    0.934    0.934 arrow_dataset.py:2916(map)
        4    0.000    0.000    0.933    0.233 arrow_dataset.py:3277(_map_single)
        1    0.000    0.000    0.883    0.883 arrow_writer.py:466(write)
        2    0.000    0.000    0.883    0.441 arrow_writer.py:424(write_examples_on_file)
        1    0.000    0.000    0.882    0.882 arrow_writer.py:528(write_batch)
        2    0.000    0.000    0.877    0.439 arrow_writer.py:161(__arrow_array__)
      4/2    0.000    0.000    0.877    0.439 table.py:1800(wrapper)
        2    0.000    0.000    0.877    0.439 table.py:1950(cast_array_to_feature)
        2    0.009    0.005    0.877    0.439 image.py:209(cast_storage)
        2    0.000    0.000    0.868    0.434 image.py:335(encode_np_array)
        2    0.000    0.000    0.856    0.428 image.py:317(image_to_bytes)
        2    0.000    0.000    0.822    0.411 Image.py:2376(save)
        2    0.000    0.000    0.822    0.411 PngImagePlugin.py:1233(_save)
        2    0.000    0.000    0.822    0.411 ImageFile.py:517(_save)
        2    0.000    0.000    0.821    0.411 ImageFile.py:545(_encode_tile)
      589    0.803    0.001    0.803    0.001 {method 'encode' of 'ImagingEncoder' objects}
```

This is of course only a test as it passes through all `numpy` arrays irrespective of if they should be an image.
Also I guess `cast_storage` is meant for casting `pyarrow` storage exclusively. 
Converting to `pyarrow` array seems like a good solution as it also handles `pytorch` tensors etc., maybe there is a more efficient way to create a PIL image from a `pyarrow` array?

Not sure how this should be handled but I would be happy to help if there is a good solution.

### Environment info

- `datasets` version: 2.18.1.dev0
- Platform: Linux-6.7.11-200.fc39.x86_64-x86_64-with-glibc2.38
- Python version: 3.12.2
- `huggingface_hub` version: 0.22.2
- PyArrow version: 15.0.2
- Pandas version: 2.2.1
- `fsspec` version: 2024.3.1