Codec#
A differentiable JPEG codec, useful to simulate compression artifacts inside a training loop.
- kornia.enhance.jpeg_codec_differentiable(input, jpeg_quality, quantization_table_y=None, quantization_table_c=None)[source]#
Differentiable JPEG encoding-decoding module.
Based on [RDPC24] [SS17], we perform differentiable JPEG encoding-decoding as follows:
\[\text{JPEG}_{\text{diff}}(I, q, QT_{y}, QT_{c}) = \hat{I}\]- Where:
\(I\) is the original image to be coded.
\(q\) is the JPEG quality controlling the compression strength.
\(QT_{y}\) is the luma quantization table.
\(QT_{c}\) is the chroma quantization table.
\(\hat{I}\) is the resulting JPEG encoded-decoded image.
- Parameters:
input (
Tensor) – the RGB image to be coded.jpeg_quality (
Tensor) – JPEG quality in the range \([0, 100]\) controlling the compression strength.quantization_table_y (
Tensor|None, optional) – quantization table for Y channel. Default: None, which will load the standard quantization table.quantization_table_c (
Tensor|None, optional) – quantization table for C channels. Default: None, which will load the standard quantization table.
- Shape:
input: \((*, 3, H, W)\).
jpeg_quality: \((1)\) or \((B)\) (if used batch dim. needs to match w/ input).
quantization_table_y: \((8, 8)\) or \((B, 8, 8)\) (if used batch dim. needs to match w/ input).
quantization_table_c: \((8, 8)\) or \((B, 8, 8)\) (if used batch dim. needs to match w/ input).
- Return type:
- Returns:
JPEG coded image of the shape \((B, 3, H, W)\)
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()
In case you want to control the quantization purly base on the quantization tables use a JPEG quality of 99.5. Setting the JPEG quality to 99.5 leads to a QT scaling of 1, see Eq. 2 of [RDPC24] for details.
>>> 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]#
Differentiable JPEG encoding-decoding module.
Based on [RDPC24] [SS17], we perform differentiable JPEG encoding-decoding as follows:
\[\text{JPEG}_{\text{diff}}(I, q, QT_{y}, QT_{c}) = \hat{I}\]- Where:
\(I\) is the original image to be coded.
\(q\) is the JPEG quality controlling the compression strength.
\(QT_{y}\) is the luma quantization table.
\(QT_{c}\) is the chroma quantization table.
\(\hat{I}\) is the resulting JPEG encoded-decoded image.
Note
The input (and output) pixel range is \([0, 1]\). In case you want to handle normalized images you are required to first perform denormalization followed by normalizing the output images again.
Note, that this implementation models the encoding-decoding mapping of JPEG in a differentiable setting, however, does not allow the excess of the JPEG-coded byte file itself. For more details please refer to [RDPC24].
This implementation is not meant for data loading. For loading JPEG images please refer to kornia.io. There we provide an optimized Rust implementation for fast JPEG loading.
- Parameters:
quantization_table_y (
Tensor|Parameter|None, optional) – quantization table for Y channel. Default: None, which will load the standard quantization table.quantization_table_c (
Tensor|Parameter|None, optional) – quantization table for C channels. Default: None, which will load the standard quantization table.
- Shape:
quantization_table_y: \((8, 8)\) or \((B, 8, 8)\) (if used batch dim. needs to match w/ image_rgb).
quantization_table_c: \((8, 8)\) or \((B, 8, 8)\) (if used batch dim. needs to match w/ image_rgb).
image_rgb: \((*, 3, H, W)\).
jpeg_quality: \((1)\) or \((B)\) (if used batch dim. needs to match w/ image_rgb).
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()