initial jpegxl zarr implementation

[fcollman]

@william-silversmith I saw your jpegxl improvements, and I know that there are several people who want to use jpegxl encoded zarr data sources. I tried to take the code you have and create an adapter. I had to make some changes to make it work for chunk decoding, and i was confused by it was in the sliceview decoder. Perhaps you could shed some light on that and provide some good feedback to this PR.

I've validated that this works with a uint16 jpegxl data volume in zarr2 as 1x1x32x256x256 chunks. I am aware that I don't think i've fully generalized this for all the use cases that people might imagine, such as 1x3x32x245x256 and storing those 3 as an RGB image.

Also if the zarr ordering is very different this will fail as it assumes the spatial dimensions are last, which i think for zarr3 is no longer mandated. Not sure how to use the existing structure to make that work well, again feedback would be welcome here.

[jbms]

For zarr3, the codec should get registered at https://github.com/zarr-developers/zarr-extensions

Rather than have complicated rules for determining the mapping from array to image dimensions, I think it would be better to have the image codecs always use a simple mapping, like require 2d [y, x] or 3d [y, x, c] input array, and then chain after transpose and/or reshape codec to allow arbitrary mapping.

My reshape proposal is here:
zarr-developers/zarr-extensions#10

[fcollman]

just going to make sure i understand. Your proposal would be that someone who has a zarr that says it is chunked at 1 x 1 x 32 x 256 x 256 would then specify a reshape codec that would take the 32 x 256 x 256 data that is natively encoded in the chunk, and then turn it into 1 x 1 x 32 x256 x256. In this way the ideal decoder wouldn't require any size input and would just spit out a 3d or 2d array, and it would be up to the rest of the codecs in the zarr chain to make sure it ends up shaped in the way that zarr spec expects?

[jbms]

just going to make sure i understand. Your proposal would be that someone who has a zarr that says it is chunked at 1 x 1 x 32 x 256 x 256 would then specify a reshape codec that would take the 32 x 256 x 256 data that is natively encoded in the chunk, and then turn it into 1 x 1 x 32 x256 x256. In this way the ideal decoder wouldn't require any size input and would just spit out a 3d or 2d array, and it would be up to the rest of the codecs in the zarr chain to make sure it ends up shaped in the way that zarr spec expects?

Yes exactly, though it would still confirm that the image has the expected size.

[fcollman]

So is this code (https://github.com/cgohlke/imagecodecs/blob/master/imagecodecs/_jpegxl.pyx#L700) the numcodecs implementation of the jpegxl decompression codec consistent with this view? I think that it is, it has some opinionated ways about how a fixed size and ordering of pixels will come out of the codec, and then it would be up to reshape/transpose to transform those things into the chunk shape advertised by the zarr metadata.

[william-silversmith]

I'm a bit sleepy so I haven't looked at everything in detail, but one thing that immediately jumped out at me was that the frames function is doing a full decode of each frame and only returning the count.

[jbms]

Regarding the imagecodecs implementation:

First, the actual numcodecs-compatible codec is defined here: https://github.com/cgohlke/imagecodecs/blob/e64d6c16eb49b006dc6116cac12997031efcf151/imagecodecs/numcodecs.py#L1125

It does include a squeeze parameter, which I am proposing should instead be handled by the separate reshape codec.

In general there seem to be various inconsistencies --- for example, the Jpeg codec (https://github.com/cgohlke/imagecodecs/blob/e64d6c16eb49b006dc6116cac12997031efcf151/imagecodecs/numcodecs.py#L949) supports some color space options while the Jpegxl codec does not.

It may be reasonable to keep the options and behavior minimal and add more options in the future (in a backwards-compatible way) as needed.

Note that these image formats have quite a lot of complexity regarding color space representation. For grayscale/single-channel images that is probably not too relevant but for multi-channel images it is important to understand the impact of color space conversions. I suspect for many microscopy applications where human perception of color is not particularly relevant, you would want to avoid any color space conversion and just compress each channel independently, which can be done but is typically not the default. Alternatively the chunking could be such that each channel is a separate chunk and then the image codec only seems grayscale images anyway.

[fcollman]

okay.. i'm going to close this PR for now, also because I can see Will's improvements are merged now and really the number of frames should follow his pattern as he suggested (one function call to read the height, width and number of frames without reading the image(s)).

also i named by branch wrong.. embarrassing, and the codec situation is clearly not resolved enough to push through a particular implementation right now, at least for v3.