Added support for JXL images.
Opened by alexjc (open)

JPEG-XL is the most advanced of the next-generation of image codecs, supporting both lossless and lossy files — with better compression and quality than PNG and JPG respectively.  It has reduced the disk sizes and bandwidth required for many of the datasets I use.

Pillow does not yet support JXL, but there's a plugin as a separate Python library that does (`pip install jxlpy`), and I've tested that this change works as expected when the plugin is imported.

Dataset used for testing, you must `git pull` as loading it from Python won't work until `datasets-server` is also changed to support JXL files:
https://huggingface.co/datasets/texturedesign/td01_natural-ground-textures

The case where the plugin is not imported first raises an error:
```
PIL.UnidentifiedImageError: cannot identify image file 'td01/train/set01/01_145523.jxl'
```

In order to enable support for JXL even before pillow supports this, should this exception be handled with a better error message?  I'd expect/hope JXL support to follow in one of the pillow quarterly releases in the next 6-9 months.

Comments:
I'm fine with the addition of jxl in the list of known image extensions, this way users that have the plugin can work with their JXL datasets. WDYT @mariosasko ?

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/datasets/pr_5289). All of your documentation changes will be reflected on that endpoint.

I think we should wait for official support from Pillow. Plus, the linked plugin doesn't support `Image.save`, which is one of the requirements for a format to be included in `IMAGE_EXTENSIONS`.

@alexjc In the meantime, one option is to add these lines to the card:
```python
import importlib
import datasets

if ".jxl" not in datasets.packaged_modules.imagefolder.IMAGE_EXTENSIONS:
  datasets.packaged_modules.imagefolder.IMAGE_EXTENSIONS.append(".jxl")

if "jxl" not in datasets.packaged_modules._EXTENSION_TO_MODULE:
  datasets.packaged_modules._EXTENSION_TO_MODULE["jxl"] = ("imagefolder", {})

importlib.reload(datasets.load)
ds = datasets.load_dataset("texturedesign/td01_natural-ground-textures")
```
And you can add a note to the card that this dataset requires the "jxlpy" package to work. 

In this case, you can also disable the viewer to avoid the discrepancy between the data displayed in the preview and the loaded data.

Another option is to define the loading script and add `jxlpy` to the list of dependencies [here](https://github.com/huggingface/datasets-server/blob/3012da62054a025467616abc14b0b46e1f11ea13/workers/first_rows/pyproject.toml#L8) to enable the viewer. This option requires more work, so let us know if you need help.

Thank you both for your thoughtful replies!

One questions and and update:
* The jxlpy plugin does support saving, in the `_save` function of the JXLImagePlugin file.  Did it not work?  I'm working on the upgrade to the latest JXL, so it'd be good to know if it failed so I can fix it.
* I wrote to the Pillow maintainer and the preferred solution would be to keep JXL as a separate plugin because they're a small team don't have the resources to maintain more code.

With that in mind, let me share the minimal set of features I'd need for this to work within the `datasets` library:
1. Using `load_dataset()` with the HuggingFace dataset name correctly downloads the JXL files so they are available locally. Even if the `file_name` field is left intact and not loaded as a PIL image, this is the first step.
2. With minimal monkey-patching, having the `load_dataset` correctly expand `file_name` into PIL `image` fields if JXL support is available.

If both of these work, then I can use HuggingFace's hub and the `datasets` library for an MVP even if not all features are there.  I don't need automatic thumbnails or previews of the dataset on the server.


Given the reply from the Pillow maintainer, what solution can we come up with that works in a more permanent way than waiting for Pillow integration (which may not happen) — assuming users install the `jxlpy` plugin separately?

Link to my upgrade for the latest `libjxl`, pending review and merge. I tested load/save via Pillow extensively for this: https://github.com/olokelo/jxlpy/pull/13

After more research, here's my latest suggestion:
* Depending on the build of pillow, the source (pip or conda), the platform even, certain formats may or may not be available — despite them being in the list.  For example, webp support is not consistently available.
* I'd suggest adding JXL to the list and simply catching the `PIL.UnidentifiedImageError` — printing a useful error message that sends them to a Wiki page to find out what to do.
* On that page would be included instructions how to install support for the format and what to do for the dataset to load correctly on any platform, both with or without conda, etc.

What do you think?

> The jxlpy plugin does support saving, in the _save function of the JXLImagePlugin file. Did it not work? I'm working on the upgrade to the latest JXL, so it'd be good to know if it failed so I can fix it.

My bad, I was referring to [this](https://github.com/google/brunsli/blob/2dd949e53ed05796eb44a31cc759fbf9e6c53e2f/contrib/py/jxl_library_patches/jxl_pillow.py) version of the plugin.

I still think this involves too much work:
* would require a new doc page
* unofficial plugins have to be imported explicitly, leading to messier code on our side
* etc.

For now, it seems more reasonable to create a loading script (faster than ImageFolder, as ImageFolder has to resolve the image files first) for this particular case and add `jxlpy` to the list of the `datasets-server`'s dependencies. Also, one additional advantage of this approach is that it reports if any of the modules imported in a script is missing, which is handy in your case for the plugin lib. WDYT?

OK, let me try it it and I'll report back.

Will the JXL files (even if unknown format) be automatically downloaded if they are linked from the `.jsonl` file?

(I had trouble getting that working before this patch.)

> Will the JXL files (even if unknown format) be automatically downloaded if they are linked from the .jsonl file?

No, they need to be downloaded explicitly.

Feel free to use 🤗 Hub discussions in your dataset repo to ping us for help (our usernames are the same there)

Is it possible to add support for JXL files being downloaded without needing to add server-side rendering support?

In the loading script, data files are downloaded with `DownloadManager` (`dl_manager` in `_split_generators`), which doesn't have any requirements regarding the actual type of the downloaded files.

PS: Let's use the forum or Hub discussions for further questions to avoid pinging other participants