Add Image feature
Opened by mariosasko (closed)

Adds the Image feature. This feature is heavily inspired by the recently added Audio feature (#2324). Currently, this PR is pretty simple.

Some considerations that need further discussion:
* I've decided to use `Pillow`/`PIL` as the image decoding library. Another candidate I considered is `torchvision`, mostly because of its `accimage` backend, which should be faster for loading `jpeg` images than `Pillow`. However, `torchvision`'s io module only supports png and jpeg images, has `torch` as a hard dependency, and requires magic to work with image bytes ( `torch.ByteTensor(torch.ByteStorage.from_buffer(image_bytes)))`).
* Currently, I'm converting `PIL`'s `Image` type to `np.ndarray`. The vision models in Transformers such as ViT prefer the raw `Image` type and not the decoded tensors, so there is a small overhead due to [this conversion](https://github.com/huggingface/transformers/blob/3e8761ab8077e3bb243fe2f78b2a682bd2257cf1/src/transformers/image_utils.py#L62-L73). IMO this is justified to keep this part aligned with the Audio feature, which also returns `np.ndarray`. What do you think?
* Still have to work on the channel decoding logic:
    * PyTorch prefers the channel-first ordering (C, H, W); TF and Flax the channel-last ordering (H, W, C). One cool feature would be adjusting the channel order based on the selected formatter (`torch`, `tf`, `jax`). 
    * By default, `Image.open` returns images of shape (H, W, C). However, ViT's feature extractor expects the format (C, H, W) if the image is passed as an array (explained [here](https://huggingface.co/transformers/model_doc/vit.html#transformers.ViTFeatureExtractor.__call__)), so I'm more inclined to the format (C, H, W). Which one do you prefer, (C, H, W) or (H, W, C)?
* Are there any options you'd like to see? (the user could change those via `cast_column`, such as `sampling_rate` in the Audio feature)


TODOs:
* [x] tests
* in subsequent PRs:
  * docs - a section in the docs, which gives some additional info on the Image and Audio feature and compares them to 
  `ArrayND`  
  * streaming (waiting for #3129 and #3133 to get merged first)
  * update the image tasks and the datasets to use the new feature
  * Image/Audio formatting

[Colab Notebook](https://colab.research.google.com/drive/1mIrTnqTVkWLJWoBzT1ABSe-LFelIep1c?usp=sharing) where you can play with this feature.

I'm also adding a link to the [Image](https://github.com/tensorflow/datasets/blob/7ac7d506488d46038a5854961d068926b3f93c7f/tensorflow_datasets/core/features/image_feature.py#L155) feature in TFDS because one of our goals is to parse TFDS scripts eventually, so our Image feature has to (at least) support all the formats theirs does.
Feel free to cc anyone who might be interested.

P.S. Please ignore the changes in the `datasets/**/*.py` files 😄.

Comments:
Awesome, looking forward to using it :)

Few additional comments:
* the current API doesn't meet the requirements mentioned in #3145 (e.g. image mime-type). However, this will be doable soon as we also plan to store image bytes alongside paths in arrow files (see https://github.com/huggingface/datasets/pull/3129#discussion_r738426187). Then, PIL can return the correct mime-type: 
   ```python
   from PIL import Image
   import io

   mimetype = Image.open(io.BytesIO(image_bytes)).get_format_mimetype()
   ``` 
   I plan to add this change in a separate PR.
*  currently, I'm returning an `np.ndarray` object after decoding for consistency with the Audio feature. However, the vision models from Transformers prefer an `Image` object to avoid the `Image.fromarray` call in the corresponding feature extractors (see [this warning](https://huggingface.co/transformers/master/model_doc/vit.html#transformers.ViTFeatureExtractor.__call__) in the Transformers docs) cc @NielsRogge 

So I'm not entirely sure whether to return only a NumPy array, only a PIL Image, or both when decoding. The last point worries me because we shouldn't provide an API that leads to a warning in Transformers (in the docs, not in code :)). At the same time, it makes sense to preserve consistency with the Audio feature and return a NumPy array. 

That's why I would appreciate your opinions on this.

That is a good question. Also pinging @nateraw .

Currently we only support returning numpy arrays because of numpy/tf/torch/jax formatting features that we have, and to keep things simple. See the [set_format docs](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasets.Dataset.set_format) for more info

I don't think centering the discussion on what ViT expects is good, as the vision Transformers model are still in an experimental stage and we can adapt those depending on what you do here :-).

IMO, the discussion should revolve on what a user will want to do with a vision dataset, and they will want to:
- lazily decode their images
- maybe apply data augmentation (for the training set)
- resize to a fixed shape for batching

The libraries that provide step 2 and 3 either use PIL (thinking torchvision) or cv2 (thinking albumentations). NumPy does not have any function to resize an image or do basic data augmentation (like a rotate) so I think it shouldn't be the default format for an image dataset, PIL or cv2 (in an ideal world with the ability to switch between the two depending on what the users prefer) would be better.

Side note: I will work on the vision integration in Transformers with Niels next month so please keep me in the loop for those awesome new vision features!

@sgugger I completely agree with you, especially after trying to convert the `run_image_classification` script from Transformers to use this feature. The current API doesn't seem intuitive there due to the torchvision transforms, which, as you say, prefer PIL over NumPy arrays. 

So the default format would return `Image` (PIL) / `np.ndarray` (cv2) and `set_format(numpy/tf/pt)` would return image tensors if I understand you correctly. IMO this makes a lot more sense (and flexibility) than the current API.

Also, one additional library worth mentioning here is AugLy which supports image file paths and `PIL.Image.Image` objects.

That's so nice !

Also I couldn't help myself so I've played with it already ^^
I was agreeably surprised that with minor additions I managed to even allow this, which I find very satisfactory:
```python
import PIL.Image
from datasets import Dataset

path = "docs/source/imgs/datasets_logo_name.jpg"

dataset = Dataset.from_dict({"img": [PIL.Image.open(path)]})
print(dataset.features)
# {'img': Image(id=None)}
print(dataset[0]["img"])
# <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=1200x300 at 0x129DE4AC8>
```

Let me know if that's a behavior you'd also like to see 

EDIT: just pushed my changes on a branch, you can see the diff [here](https://github.com/mariosasko/datasets-1/compare/add-image-feature...huggingface:image-type-inference) if you want

Thanks, @lhoestq! I like your change. Very elegant indeed.

P.S. I have to write a big comment that explains all the changes/things left to consider. Will do that in the next few days!

I'm marking this PR as ready for review.

Thanks to @sgugger's comment, the API is much more flexible now as it decodes images (lazily) as `PIL.Image.Image` objects and supports transforms directly on them.

Also, we no longer return paths explicitly (previously, we would return `{"path": image_path, "image": pil_image}`) for the following reasons:
* what to return when reading an image from an URL or a NumPy array. We could set `path` to `None` in these situations, but IMO we should avoid redundant information.
* returning a dict doesn't match nicely with the requirement of supporting image modifications - what to do if the user modifies both the image path and the image

(Btw, for the images stored locally, you can access their paths with `dset[idx]["image"].filename`, or by avoiding decoding with `paths = [ex["path"] for ex in dset]`. @lhoestq @albertvillanova WDYT about having an option to skip decoding for complex features, e. g. `Audio(decode=False)`? This way, the user can easily access the underlying data.)

Examples of what you can do:
```python
# load local images
dset = Dataset.from_dict("image": [local_image_path], features=Features({"images": Image()}))
# load remote images (we got this for free by adding support for streaming)
dset = Dataset.from_dict("image": [image_url], features=Features({"images": Image()}))
# from np.ndarray
dset = Dataset.from_dict({"image": [np.array(...)]}, features=Features({"images": Image()}))
# cast column
dset = Dataset.from_dict({"image": [local_image_path]})
dset.cast_column("image", Image())

# automatic type inference
dset = Dataset.from_dict({"image": [PIL.Image.open(local_image_path)]})

# transforms
def img_transform(example):
     ...
     example["image"] = transformed_pil_image_or_np_ndarray
     return example
dset.map(img_trnasform)

# transform that adds a new column with images (automatic inference of the feature type)
dset.map(lambda ex: {"image_resized": ex["image"].resize((100, 100))})
print(dset.features["image_resized"]) # will print Image()
```

Some more cool features:
* We store the image filename (`pil_image.filename`) whenever possible to avoid costly conversion to bytes
* if possible, we use native compression when encoding images. Otherwise, we fall back to the lossless PNG format (e.g. after image ops or when storing NumPy arrays)

Hints to make reviewing easier:
* feel free to ignore the extension type part because it's related to PyArrow internals.
* also, let me know if we are too strict/ too flexible in terms of types the Image feature can encode/decode. Hints:
   * `encode_example` handles encoding during dataset generation (you can think of it as `yield key, features.encode_example(example)`)
   * `objects_to_list_of_image_dicts` handles encoding of returned examples in `map`

P.S. I'll fork the PR branch and start adding the Image feature to the existing image datasets (will also update the `ImageClassification` template while doing that).

> WDYT about having an option to skip decoding for complex features, e. g. Audio(decode=False)?

Yes definitely, also I think it could be useful for the dataset viewer to not decode the data but instead return either the bytes or the (possibly chained) URL. cc @severo 

We want to merge this today/tomorrow, so I'd really appreciate your reviews @sgugger @nateraw.

Also, you can test this feature on the existing image datasets (MNIST, beans, food101, ...) by installing `datasets` from the PR branch:
```
pip install git+https://github.com/huggingface/datasets.git@adapt-image-datasets
```


Thanks for the review @nateraw!

1. This is a copy of your notebook with the fixed map call: https://colab.research.google.com/gist/mariosasko/e351a717682a9392ca03908e65a2600e/image-feature-demo.ipynb
   (Sorry for misleading you with the map call in my un-updated notebook)
   Also, we can avoid this cast by trying to infer the type of the column (`"pixel_values"`) returned by the image feature extractor (we are already doing something similar for the columns with names: `"attention_mask"`, `"input_ids"`, ...). I plan to add this QOL improvement soon. 
2. It should work OK even without updating Pillow and PyArrow (these two libraries are pre-installed in Colab, so updating them requires a restart of the runtime). 
     > I noticed an error that I'm guessing you ran into when I tried using the older version

    Do you recall which type of error it was because everything works fine on my side if I run the notebooks with the lowest supported version of Pillow (`6.2.1`)?

Thanks for playing with it @nateraw and for sharing your notebook, this is useful :)

I think this is ready now, congrats @mariosasko !

Love this feature and hope to release soon!