torchjpeg.quantization

General Quantization

The torchjpeg.quantization package provides functions which quantize DCT coefficients. The IJG (libjpeg) quantization matrices are included as part of this package as well as code which generates them from a scalar quality factor. Users can also provide their own quantization matrices. This implementation of the IJG quantization matrices only deals with the “baseline” setting where the maximum quantization value is 255. Functions in this module operate on single channel images since the channels are often quantized separately or may not be at the same resolution.

torchjpeg.quantization.quantize(dct: torch.Tensor, mat: torch.Tensor, round_func: Callable[[torch.Tensor], torch.Tensor] = <built-in method round of type object>) → torch.Tensor[source]

Quantizes DCT coefficients.

Parameters
  • dct (Tensor) – DCT coefficients of shape \((N, 1, H, W)\).

  • mat – (Tensor): Quantization matrix of shape \((1, 1, 8, 8)\).

  • round_func – (Callable[[Tensor], Tensor]): Rounding function to use, defaults to torch.round().

Returns

Quantized DCT coefficients.

Return type

Tensor

Note

DCT quantization is computed as

\[\widetilde{D}_{ij} = \left\lfloor \frac{D_{ij}}{Q_{ij}} \right\rceil \]

For DCT coefficients \(D\) and quantization matrix \(Q\).

torchjpeg.quantization.dequantize(dct: torch.Tensor, mat: torch.Tensor) → torch.Tensor[source]

Dequantize DCT coefficients.

Parameters
  • dct (Tensor) – Quantized DCT coefficients of shape \((N, 1, H, W)\).

  • mat – (Tensor): Quantization matrix of shape \((1, 1, 8, 8)\).

Returns

Quantized DCT coefficients.

Return type

Tensor

Note

DCT dequantization is computed as

\[D_{ij} = \widetilde{D}_{ij} \cdot Q_{ij} \]

For quantized DCT coefficients \(\widetilde{D}\) and quantization matrix \(Q\).

torchjpeg.quantization.quantize_multichannel()[source]

Quantizes a three channel image of DCT coefficients.

Parameters
  • dct (Tensor) – DCT coefficients of shape \((N, 3, H, W)\).

  • mat (Tensor) – Quantization matrix of shape \((1, 2, 8, 8)\).

  • round – (Callable[[Tensor], Tensor]): Rounding function to use, defaults to torch.round().

Returns

  • Tensor – Y channel coefficients of shape \((N, 1, H, W)\).

  • Tensor – Cb channel coefficients of shape \(\left(N, 1, \frac{H}{2}, \frac{W}{2}\right)\).

  • Tensor – Cr channel coefficients of shape \(\left(N, 1, \frac{H}{2}, \frac{W}{2}\right)\).

Note

This function performs chroma subsampling

torchjpeg.quantization.dequantize_multichannel(y: torch.Tensor, cb: torch.Tensor, cr: torch.Tensor, mat: torch.Tensor) → torch.Tensor[source]

Dequantizes a three channel image.

Parameters
  • y (Tensor) – Quantized Y channel coefficients of shape \((N, 1, H, W)\).

  • cb (Tensor) – Quantized Cb channel coefficients of shape \(\left(N, 1, \frac{H}{2}, \frac{W}{2}\right)\).

  • cr (Tensor) – Quantized Cr channel coefficients of shape \(\left(N, 1, \frac{H}{2}, \frac{W}{2}\right)\).

Returns

A three channel image of DCT coefficients.

Return type

Tensor

Note

This function assumes chroma subsampling.

IJG (libjpeg) Compatible Quantization

torchjpeg.quantization.ijg provides functions which match the Independent JPEG Group’s libjpeg quantization method.

torchjpeg.quantization.ijg.compress_coefficients(batch: torch.Tensor, quality: int, table: str = 'luma') → torch.Tensor[source]

A high level function that takes a batch of pixels in [0, 1] and returns quantized DCT coefficients.

Parameters
  • batch (Tensor) – A batch of images to quantize of shape (N, 1, H, W).

  • quality (int) – A scalar quality in [0, 100] specifying the desired quality.

  • table (str) – One of “luma” or “chroma” to choose the desired set of tables.

Returns

A batch of quantized DCT coefficients.

Return type

Tensor

torchjpeg.quantization.ijg.decompress_coefficients(batch: torch.Tensor, quality: int, table: str = 'luma') → torch.Tensor[source]

A high level function that converts quantized DCT coefficients to pixels.

Parameters
  • batch (Tensor) – A batch of quantized DCT coefficients of shape (N, 1, H, W).

  • quality (int) – A scalar quality in [0, 100] specifying the quality the coefficients were quantized at.

  • table (str) – One of “luma” or “chroma” to choose the desired set of tables.

Returns

A batch of image pixels.

Return type

Tensor

torchjpeg.quantization.ijg.dequantize_at_quality(dct: torch.Tensor, quality: int, table: str = 'luma') → torch.Tensor[source]

Dequantizes using a scalar quality instead of a quantization matrix. uses IJG quantization matrices.

Parameters
  • dct (Tensor) – Quantized DCT coefficients of shape \((N, 1, H, W)\).

  • quality (int) – A scalar in [0, 100] specifying the quality that the coefficients were quantized at.

  • table (str) – One of “luma” or “chroma” to choose the desired set of tables.

torchjpeg.quantization.ijg.get_coefficients_for_qualities(quality: torch.Tensor, table: str = 'luma') → torch.Tensor[source]

Gets IJG quantization matrices for a given batch of qualities.

Parameters
  • quality (Tensor) – A batch of qualities of shape \((N)\)

  • table (str) – A string indicating the table to use, either “luma” or “chroma”

torchjpeg.quantization.ijg.qualities_to_scale_factors(qualities: torch.Tensor) → torch.Tensor[source]

Converts a batch of qualities in [0, 100] to a batch of scale factors suitable for scaling one of the IJG reference quantization matrices.

Parameters

qualities (Tensor) – A single dimensional batch of qualities.

Returns

A single dimensional batch of scale factors.

Return type

Tensor

torchjpeg.quantization.ijg.quantize_at_quality(dct: torch.Tensor, quality: int, table: str = 'luma') → torch.Tensor[source]

Quantizes using a scalar quality instead of a quantization matrix. Uses IJG quantization matrices.

Parameters
  • dct (Tensor) – DCT coefficients of shape \((N, 1, H ,W)\).

  • quality (int) – A scalar in [0, 100] specifying the desired quality.

  • table (str) – One of “luma” or “chroma” to choose the desired set of tables.

Returns

Quantized DCT coefficients.

Return type

Tensor

torchjpeg.quantization.ijg.scale_quantization_matrices(scale_factor: torch.Tensor, table: str = 'luma') → torch.Tensor[source]

Scales one of the IJG reference quantization matrices.

Parameters
  • scale_factor (Tensor) – A batch of \(N\) scale factors.

  • table (str) – A string indicating the table to use, either “luma” or “chroma”

Returns

A batch of quantization matrices of shape \((N, 64)\).

Return type

Tensor