Codec#

A differentiable JPEG codec, to simulate compression artifacts inside a training loop. An integer image raises TypeCheckError; convert 8-bit data with image.float() / 255 first.

kornia.enhance.jpeg_codec_differentiable(input, jpeg_quality, quantization_table_y=None, quantization_table_c=None)[source]#

Compress and decompress an RGB image with a differentiable JPEG model.

Try in browser

out = jpeg_codec_differentiable(input, jpeg_quality) returns an RGB image in \([0, 1]\) after a JPEG encode and decode at jpeg_quality, following [RDPC24] and [SS17]. Gradients reach input, jpeg_quality and the quantization tables.

Convention:

Input and output pixels are in \([0, 1]\): denormalize a normalized image first and normalize the output again. Any H and W work: the codec pads the bottom and right to a multiple of 16 and crops back.
The decoded image is clamped to \([0, 1]\) by ycbcr_to_rgb(), so a pixel the codec pushes past either end passes no gradient.
The result approximates libjpeg and does not reproduce its output exactly.
float16 and bfloat16 inputs are computed in float32, quality and tables included, and the output is rounded once to the input dtype.
The codec runs with autocast disabled, so inside a torch.autocast region it gives the same output as outside it, and the same gradients when the backward pass runs after the region (on CUDA, up to the run-to-run spread of its atomic accumulation unless deterministic algorithms are enabled). A backward pass inside the region, which torch.compile assumes by default, still computes part of the gradient in the autocast dtype.

Parameters:

input (Tensor) – RGB image with shape (*, 3, H, W).
jpeg_quality (Tensor) – JPEG quality in \([0, 100]\), with shape (1,) for all images or (N,) for N the number of flattened images; a 0-d tensor raises. A quality of exactly 0 is treated as 1, as in libjpeg; a quality between 0 and 1 is left alone and compresses harder still.
quantization_table_y (Tensor | None, optional) – luma quantization table with shape (8, 8) for all images or (N, 8, 8). The default is the standard table. Default: None
quantization_table_c (Tensor | None, optional) – chroma quantization table, in the same shapes. The default is the standard table. Default: None

Return type:

Tensor

Returns:

JPEG-coded image with the shape of input.

Note

The model produces no JPEG byte file and cannot read one. To load JPEG files use kornia.io.

See also

JPEGCodecDifferentiable: the module form, which keeps the tables.
RandomJPEG: applies JPEG compression at random quality per image.

Example

To perform JPEG coding with the standard quantization tables just provide a JPEG quality

>>> img = torch.rand(3, 3, 64, 64, requires_grad=True, dtype=torch.float)
>>> jpeg_quality = torch.tensor((99.0, 25.0, 1.0), requires_grad=True)
>>> img_jpeg = jpeg_codec_differentiable(img, jpeg_quality)
>>> img_jpeg.sum().backward()

You also have the option to provide custom quantization tables

>>> img = torch.rand(3, 3, 64, 64, requires_grad=True, dtype=torch.float)
>>> jpeg_quality = torch.tensor((99.0, 25.0, 1.0), requires_grad=True)
>>> quantization_table_y = torch.randint(1, 256, size=(3, 8, 8), dtype=torch.float)
>>> quantization_table_c = torch.randint(1, 256, size=(3, 8, 8), dtype=torch.float)
>>> img_jpeg = jpeg_codec_differentiable(img, jpeg_quality, quantization_table_y, quantization_table_c)
>>> img_jpeg.sum().backward()

To control the quantization purely with the quantization tables use a JPEG quality of 99.5, which scales the tables by 1 (Eq. 2 of [RDPC24]).

>>> img = torch.rand(3, 3, 64, 64, requires_grad=True, dtype=torch.float)
>>> jpeg_quality = torch.ones(3) * 99.5
>>> quantization_table_y = torch.randint(1, 256, size=(3, 8, 8), dtype=torch.float)
>>> quantization_table_c = torch.randint(1, 256, size=(3, 8, 8), dtype=torch.float)
>>> img_jpeg = jpeg_codec_differentiable(img, jpeg_quality, quantization_table_y, quantization_table_c)
>>> img_jpeg.sum().backward()

class kornia.enhance.JPEGCodecDifferentiable(quantization_table_y=None, quantization_table_c=None)[source]#

Compress and decompress an RGB image with a differentiable JPEG model (module form).

Try in browser

op = JPEGCodecDifferentiable(quantization_table_y, quantization_table_c) stores the tables, and out = op(image_rgb, jpeg_quality) is jpeg_codec_differentiable() with them. A table is stored as a buffer, or as a learnable parameter if it is an nn.Parameter.

Convention:

See the convention block of jpeg_codec_differentiable().

Parameters:

quantization_table_y (Tensor | Parameter | None, optional) – luma quantization table with shape (8, 8) for all images or (N, 8, 8). The default is the standard table. Default: None
quantization_table_c (Tensor | Parameter | None, optional) – chroma quantization table, in the same shapes. The default is the standard table. Default: None

Shape:

image_rgb: (*, 3, H, W).
jpeg_quality: (1,) or (N,) for N the number of flattened images.
Output: same shape as image_rgb.

See also

jpeg_codec_differentiable(): the one-call form.

Example

You can use the differentiable JPEG module with standard quantization tables by

>>> diff_jpeg_module = JPEGCodecDifferentiable()
>>> img = torch.rand(2, 3, 32, 32, requires_grad=True, dtype=torch.float)
>>> jpeg_quality = torch.tensor((99.0, 1.0), requires_grad=True)
>>> img_jpeg = diff_jpeg_module(img, jpeg_quality)
>>> img_jpeg.sum().backward()

You can also specify custom quantization tables to be used by

>>> quantization_table_y = torch.randint(1, 256, size=(2, 8, 8), dtype=torch.float)
>>> quantization_table_c = torch.randint(1, 256, size=(2, 8, 8), dtype=torch.float)
>>> diff_jpeg_module = JPEGCodecDifferentiable(quantization_table_y, quantization_table_c)
>>> img = torch.rand(2, 3, 32, 32, requires_grad=True, dtype=torch.float)
>>> jpeg_quality = torch.tensor((99.0, 1.0), requires_grad=True)
>>> img_jpeg = diff_jpeg_module(img, jpeg_quality)
>>> img_jpeg.sum().backward()

In case you want to learn the quantization tables just pass parameters nn.Parameter

>>> quantization_table_y = torch.nn.Parameter(torch.randint(1, 256, size=(2, 8, 8), dtype=torch.float))
>>> quantization_table_c = torch.nn.Parameter(torch.randint(1, 256, size=(2, 8, 8), dtype=torch.float))
>>> diff_jpeg_module = JPEGCodecDifferentiable(quantization_table_y, quantization_table_c)
>>> img = torch.rand(2, 3, 32, 32, requires_grad=True, dtype=torch.float)
>>> jpeg_quality = torch.tensor((99.0, 1.0), requires_grad=True)
>>> img_jpeg = diff_jpeg_module(img, jpeg_quality)
>>> img_jpeg.sum().backward()