Support LPIPS when using LAB colorspace.

[pcuenca]

Addresses #12.

I also double-checked that the input for LPIPS are images in the range [-1, 1], everything looks correct in that regard.

[borisdayma]

Can you call the function to_rgb (or similar) and make it a noop when the format is already RGB?

[borisdayma]

Also can you check if it affects the prepare_dataset notebook and the wandb logging (may need to use a jax.device_get there).

[pcuenca]

Can you call the function to_rgb (or similar) and make it a noop when the format is already RGB?

I didn't want to call it to_rgb because then it looks semantically similar to logits_to_image, but images are returned in [0, 1] and this one works in [-1, 1]. But maybe it's clearer that way.

[borisdayma]

Can you call the function to_rgb (or similar) and make it a noop when the format is already RGB?

I didn't want to call it to_rgb because then it looks semantically similar to logits_to_image, but images are returned in [0, 1] and this one works in [-1, 1]. But maybe it's clearer that way.

Maybe to_lpips? Pretty close to what you had then.

[pcuenca]

Also can you check if it affects the prepare_dataset notebook and the wandb logging (may need to use a jax.device_get there).

Sure, good point. But this was branched from main not stylegan. I'll check if there are any changes needed and if so maybe it makes sense to submit as a PR to stylegan instead.

[pcuenca]

I addressed a couple of @borisdayma's comments, but it turns out we can't use tfio.experimental.color.lab_to_rgb in jitted code because there's an implicit conversion to numpy here.

I checked the color conversion functions in dm_pix (https://github.com/deepmind/dm_pix/blob/master/dm_pix/_src/color_conversion.py), but they don't have LAB support yet.

We can port TensorFlow's implementation to JAX, but maybe we have more important things to do first. What do you think @borisdayma, @patil-suraj ?