Source code
Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
#include "lib/jxl/enc_quant_weights.h"
#include <jxl/memory_manager.h>
#include <jxl/types.h>
#include <cmath>
#include <cstdlib>
#include "lib/jxl/base/common.h"
#include "lib/jxl/base/status.h"
#include "lib/jxl/enc_aux_out.h"
#include "lib/jxl/enc_bit_writer.h"
#include "lib/jxl/enc_modular.h"
#include "lib/jxl/fields.h"
#include "lib/jxl/modular/encoding/encoding.h"
namespace jxl {
namespace {
Status EncodeDctParams(const DctQuantWeightParams& params, BitWriter* writer) {
JXL_ENSURE(params.num_distance_bands >= 1);
writer->Write(DctQuantWeightParams::kLog2MaxDistanceBands,
params.num_distance_bands - 1);
for (size_t c = 0; c < 3; c++) {
for (size_t i = 0; i < params.num_distance_bands; i++) {
JXL_RETURN_IF_ERROR(F16Coder::Write(
params.distance_bands[c][i] * (i == 0 ? (1 / 64.0f) : 1.0f), writer));
}
}
return true;
}
Status EncodeQuant(JxlMemoryManager* memory_manager,
const QuantEncoding& encoding, size_t idx, size_t size_x,
size_t size_y, BitWriter* writer,
ModularFrameEncoder* modular_frame_encoder) {
writer->Write(kLog2NumQuantModes, encoding.mode);
size_x *= kBlockDim;
size_y *= kBlockDim;
switch (encoding.mode) {
case QuantEncoding::kQuantModeLibrary: {
writer->Write(kCeilLog2NumPredefinedTables, encoding.predefined);
break;
}
case QuantEncoding::kQuantModeID: {
for (size_t c = 0; c < 3; c++) {
for (size_t i = 0; i < 3; i++) {
JXL_RETURN_IF_ERROR(
F16Coder::Write(encoding.idweights[c][i] * (1.0f / 64), writer));
}
}
break;
}
case QuantEncoding::kQuantModeDCT2: {
for (size_t c = 0; c < 3; c++) {
for (size_t i = 0; i < 6; i++) {
JXL_RETURN_IF_ERROR(F16Coder::Write(
encoding.dct2weights[c][i] * (1.0f / 64), writer));
}
}
break;
}
case QuantEncoding::kQuantModeDCT4X8: {
for (size_t c = 0; c < 3; c++) {
JXL_RETURN_IF_ERROR(
F16Coder::Write(encoding.dct4x8multipliers[c], writer));
}
JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params, writer));
break;
}
case QuantEncoding::kQuantModeDCT4: {
for (size_t c = 0; c < 3; c++) {
for (size_t i = 0; i < 2; i++) {
JXL_RETURN_IF_ERROR(
F16Coder::Write(encoding.dct4multipliers[c][i], writer));
}
}
JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params, writer));
break;
}
case QuantEncoding::kQuantModeDCT: {
JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params, writer));
break;
}
case QuantEncoding::kQuantModeRAW: {
JXL_RETURN_IF_ERROR(ModularFrameEncoder::EncodeQuantTable(
memory_manager, size_x, size_y, writer, encoding, idx,
modular_frame_encoder));
break;
}
case QuantEncoding::kQuantModeAFV: {
for (size_t c = 0; c < 3; c++) {
for (size_t i = 0; i < 9; i++) {
JXL_RETURN_IF_ERROR(F16Coder::Write(
encoding.afv_weights[c][i] * (i < 6 ? 1.0f / 64 : 1.0f), writer));
}
}
JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params, writer));
JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params_afv_4x4, writer));
break;
}
}
return true;
}
} // namespace
Status DequantMatricesEncode(JxlMemoryManager* memory_manager,
const DequantMatrices& matrices, BitWriter* writer,
LayerType layer, AuxOut* aux_out,
ModularFrameEncoder* modular_frame_encoder) {
bool all_default = true;
const std::vector<QuantEncoding>& encodings = matrices.encodings();
for (const auto& encoding : encodings) {
if (encoding.mode != QuantEncoding::kQuantModeLibrary ||
encoding.predefined != 0) {
all_default = false;
}
}
// TODO(janwas): better bound
return writer->WithMaxBits(512 * 1024, layer, aux_out, [&]() -> Status {
writer->Write(1, TO_JXL_BOOL(all_default));
if (!all_default) {
for (size_t i = 0; i < encodings.size(); i++) {
JXL_RETURN_IF_ERROR(EncodeQuant(memory_manager, encodings[i], i,
DequantMatrices::required_size_x[i],
DequantMatrices::required_size_y[i],
writer, modular_frame_encoder));
}
}
return true;
});
}
Status DequantMatricesEncodeDC(const DequantMatrices& matrices,
BitWriter* writer, LayerType layer,
AuxOut* aux_out) {
bool all_default = true;
const float* dc_quant = matrices.DCQuants();
for (size_t c = 0; c < 3; c++) {
if (dc_quant[c] != kDCQuant[c]) {
all_default = false;
}
}
return writer->WithMaxBits(
1 + sizeof(float) * kBitsPerByte * 3, layer, aux_out, [&]() -> Status {
writer->Write(1, TO_JXL_BOOL(all_default));
if (!all_default) {
for (size_t c = 0; c < 3; c++) {
JXL_RETURN_IF_ERROR(F16Coder::Write(dc_quant[c] * 128.0f, writer));
}
}
return true;
});
}
Status DequantMatricesSetCustomDC(JxlMemoryManager* memory_manager,
DequantMatrices* matrices, const float* dc) {
matrices->SetDCQuant(dc);
// Roundtrip encode/decode DC to ensure same values as decoder.
BitWriter writer{memory_manager};
// TODO(eustas): should it be LayerType::Quant?
JXL_RETURN_IF_ERROR(
DequantMatricesEncodeDC(*matrices, &writer, LayerType::Header, nullptr));
writer.ZeroPadToByte();
BitReader br(writer.GetSpan());
// Called only in the encoder: should fail only for programmer errors.
JXL_RETURN_IF_ERROR(matrices->DecodeDC(&br));
JXL_RETURN_IF_ERROR(br.Close());
return true;
}
Status DequantMatricesScaleDC(JxlMemoryManager* memory_manager,
DequantMatrices* matrices, const float scale) {
float dc[3];
for (size_t c = 0; c < 3; ++c) {
dc[c] = matrices->InvDCQuant(c) * (1.0f / scale);
}
JXL_RETURN_IF_ERROR(DequantMatricesSetCustomDC(memory_manager, matrices, dc));
return true;
}
Status DequantMatricesRoundtrip(JxlMemoryManager* memory_manager,
DequantMatrices* matrices) {
// Do not pass modular en/decoder, as they only change entropy and not
// values.
BitWriter writer{memory_manager};
// TODO(eustas): should it be LayerType::Quant?
JXL_RETURN_IF_ERROR(DequantMatricesEncode(memory_manager, *matrices, &writer,
LayerType::Header, nullptr));
writer.ZeroPadToByte();
BitReader br(writer.GetSpan());
// Called only in the encoder: should fail only for programmer errors.
JXL_RETURN_IF_ERROR(matrices->Decode(memory_manager, &br));
JXL_RETURN_IF_ERROR(br.Close());
return true;
}
Status DequantMatricesSetCustom(DequantMatrices* matrices,
const std::vector<QuantEncoding>& encodings,
ModularFrameEncoder* encoder) {
JXL_ENSURE(encoder != nullptr);
JXL_ENSURE(encodings.size() == kNumQuantTables);
JxlMemoryManager* memory_manager = encoder->memory_manager();
matrices->SetEncodings(encodings);
for (size_t i = 0; i < encodings.size(); i++) {
if (encodings[i].mode == QuantEncodingInternal::kQuantModeRAW) {
JXL_RETURN_IF_ERROR(encoder->AddQuantTable(
DequantMatrices::required_size_x[i] * kBlockDim,
DequantMatrices::required_size_y[i] * kBlockDim, encodings[i], i));
}
}
JXL_RETURN_IF_ERROR(DequantMatricesRoundtrip(memory_manager, matrices));
return true;
}
} // namespace jxl