enc_quant_weights.cc (7959B)
1 // Copyright (c) the JPEG XL Project Authors. All rights reserved. 2 // 3 // Use of this source code is governed by a BSD-style 4 // license that can be found in the LICENSE file. 5 6 #include "lib/jxl/enc_quant_weights.h" 7 8 #include <jxl/memory_manager.h> 9 #include <jxl/types.h> 10 11 #include <cmath> 12 #include <cstdlib> 13 14 #include "lib/jxl/base/common.h" 15 #include "lib/jxl/base/status.h" 16 #include "lib/jxl/enc_aux_out.h" 17 #include "lib/jxl/enc_bit_writer.h" 18 #include "lib/jxl/enc_modular.h" 19 #include "lib/jxl/fields.h" 20 #include "lib/jxl/modular/encoding/encoding.h" 21 22 namespace jxl { 23 24 namespace { 25 26 Status EncodeDctParams(const DctQuantWeightParams& params, BitWriter* writer) { 27 JXL_ENSURE(params.num_distance_bands >= 1); 28 writer->Write(DctQuantWeightParams::kLog2MaxDistanceBands, 29 params.num_distance_bands - 1); 30 for (size_t c = 0; c < 3; c++) { 31 for (size_t i = 0; i < params.num_distance_bands; i++) { 32 JXL_RETURN_IF_ERROR(F16Coder::Write( 33 params.distance_bands[c][i] * (i == 0 ? (1 / 64.0f) : 1.0f), writer)); 34 } 35 } 36 return true; 37 } 38 39 Status EncodeQuant(JxlMemoryManager* memory_manager, 40 const QuantEncoding& encoding, size_t idx, size_t size_x, 41 size_t size_y, BitWriter* writer, 42 ModularFrameEncoder* modular_frame_encoder) { 43 writer->Write(kLog2NumQuantModes, encoding.mode); 44 size_x *= kBlockDim; 45 size_y *= kBlockDim; 46 switch (encoding.mode) { 47 case QuantEncoding::kQuantModeLibrary: { 48 writer->Write(kCeilLog2NumPredefinedTables, encoding.predefined); 49 break; 50 } 51 case QuantEncoding::kQuantModeID: { 52 for (size_t c = 0; c < 3; c++) { 53 for (size_t i = 0; i < 3; i++) { 54 JXL_RETURN_IF_ERROR( 55 F16Coder::Write(encoding.idweights[c][i] * (1.0f / 64), writer)); 56 } 57 } 58 break; 59 } 60 case QuantEncoding::kQuantModeDCT2: { 61 for (size_t c = 0; c < 3; c++) { 62 for (size_t i = 0; i < 6; i++) { 63 JXL_RETURN_IF_ERROR(F16Coder::Write( 64 encoding.dct2weights[c][i] * (1.0f / 64), writer)); 65 } 66 } 67 break; 68 } 69 case QuantEncoding::kQuantModeDCT4X8: { 70 for (size_t c = 0; c < 3; c++) { 71 JXL_RETURN_IF_ERROR( 72 F16Coder::Write(encoding.dct4x8multipliers[c], writer)); 73 } 74 JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params, writer)); 75 break; 76 } 77 case QuantEncoding::kQuantModeDCT4: { 78 for (size_t c = 0; c < 3; c++) { 79 for (size_t i = 0; i < 2; i++) { 80 JXL_RETURN_IF_ERROR( 81 F16Coder::Write(encoding.dct4multipliers[c][i], writer)); 82 } 83 } 84 JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params, writer)); 85 break; 86 } 87 case QuantEncoding::kQuantModeDCT: { 88 JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params, writer)); 89 break; 90 } 91 case QuantEncoding::kQuantModeRAW: { 92 JXL_RETURN_IF_ERROR(ModularFrameEncoder::EncodeQuantTable( 93 memory_manager, size_x, size_y, writer, encoding, idx, 94 modular_frame_encoder)); 95 break; 96 } 97 case QuantEncoding::kQuantModeAFV: { 98 for (size_t c = 0; c < 3; c++) { 99 for (size_t i = 0; i < 9; i++) { 100 JXL_RETURN_IF_ERROR(F16Coder::Write( 101 encoding.afv_weights[c][i] * (i < 6 ? 1.0f / 64 : 1.0f), writer)); 102 } 103 } 104 JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params, writer)); 105 JXL_RETURN_IF_ERROR(EncodeDctParams(encoding.dct_params_afv_4x4, writer)); 106 break; 107 } 108 } 109 return true; 110 } 111 112 } // namespace 113 114 Status DequantMatricesEncode(JxlMemoryManager* memory_manager, 115 const DequantMatrices& matrices, BitWriter* writer, 116 LayerType layer, AuxOut* aux_out, 117 ModularFrameEncoder* modular_frame_encoder) { 118 bool all_default = true; 119 const std::vector<QuantEncoding>& encodings = matrices.encodings(); 120 121 for (const auto& encoding : encodings) { 122 if (encoding.mode != QuantEncoding::kQuantModeLibrary || 123 encoding.predefined != 0) { 124 all_default = false; 125 } 126 } 127 // TODO(janwas): better bound 128 return writer->WithMaxBits(512 * 1024, layer, aux_out, [&]() -> Status { 129 writer->Write(1, TO_JXL_BOOL(all_default)); 130 if (!all_default) { 131 for (size_t i = 0; i < encodings.size(); i++) { 132 JXL_RETURN_IF_ERROR(EncodeQuant(memory_manager, encodings[i], i, 133 DequantMatrices::required_size_x[i], 134 DequantMatrices::required_size_y[i], 135 writer, modular_frame_encoder)); 136 } 137 } 138 return true; 139 }); 140 } 141 142 Status DequantMatricesEncodeDC(const DequantMatrices& matrices, 143 BitWriter* writer, LayerType layer, 144 AuxOut* aux_out) { 145 bool all_default = true; 146 const float* dc_quant = matrices.DCQuants(); 147 for (size_t c = 0; c < 3; c++) { 148 if (dc_quant[c] != kDCQuant[c]) { 149 all_default = false; 150 } 151 } 152 return writer->WithMaxBits( 153 1 + sizeof(float) * kBitsPerByte * 3, layer, aux_out, [&]() -> Status { 154 writer->Write(1, TO_JXL_BOOL(all_default)); 155 if (!all_default) { 156 for (size_t c = 0; c < 3; c++) { 157 JXL_RETURN_IF_ERROR(F16Coder::Write(dc_quant[c] * 128.0f, writer)); 158 } 159 } 160 return true; 161 }); 162 } 163 164 Status DequantMatricesSetCustomDC(JxlMemoryManager* memory_manager, 165 DequantMatrices* matrices, const float* dc) { 166 matrices->SetDCQuant(dc); 167 // Roundtrip encode/decode DC to ensure same values as decoder. 168 BitWriter writer{memory_manager}; 169 // TODO(eustas): should it be LayerType::Quant? 170 JXL_RETURN_IF_ERROR( 171 DequantMatricesEncodeDC(*matrices, &writer, LayerType::Header, nullptr)); 172 writer.ZeroPadToByte(); 173 BitReader br(writer.GetSpan()); 174 // Called only in the encoder: should fail only for programmer errors. 175 JXL_RETURN_IF_ERROR(matrices->DecodeDC(&br)); 176 JXL_RETURN_IF_ERROR(br.Close()); 177 return true; 178 } 179 180 Status DequantMatricesScaleDC(JxlMemoryManager* memory_manager, 181 DequantMatrices* matrices, const float scale) { 182 float dc[3]; 183 for (size_t c = 0; c < 3; ++c) { 184 dc[c] = matrices->InvDCQuant(c) * (1.0f / scale); 185 } 186 JXL_RETURN_IF_ERROR(DequantMatricesSetCustomDC(memory_manager, matrices, dc)); 187 return true; 188 } 189 190 Status DequantMatricesRoundtrip(JxlMemoryManager* memory_manager, 191 DequantMatrices* matrices) { 192 // Do not pass modular en/decoder, as they only change entropy and not 193 // values. 194 BitWriter writer{memory_manager}; 195 // TODO(eustas): should it be LayerType::Quant? 196 JXL_RETURN_IF_ERROR(DequantMatricesEncode(memory_manager, *matrices, &writer, 197 LayerType::Header, nullptr)); 198 writer.ZeroPadToByte(); 199 BitReader br(writer.GetSpan()); 200 // Called only in the encoder: should fail only for programmer errors. 201 JXL_RETURN_IF_ERROR(matrices->Decode(memory_manager, &br)); 202 JXL_RETURN_IF_ERROR(br.Close()); 203 return true; 204 } 205 206 Status DequantMatricesSetCustom(DequantMatrices* matrices, 207 const std::vector<QuantEncoding>& encodings, 208 ModularFrameEncoder* encoder) { 209 JXL_ENSURE(encoder != nullptr); 210 JXL_ENSURE(encodings.size() == kNumQuantTables); 211 JxlMemoryManager* memory_manager = encoder->memory_manager(); 212 matrices->SetEncodings(encodings); 213 for (size_t i = 0; i < encodings.size(); i++) { 214 if (encodings[i].mode == QuantEncodingInternal::kQuantModeRAW) { 215 JXL_RETURN_IF_ERROR(encoder->AddQuantTable( 216 DequantMatrices::required_size_x[i] * kBlockDim, 217 DequantMatrices::required_size_y[i] * kBlockDim, encodings[i], i)); 218 } 219 } 220 JXL_RETURN_IF_ERROR(DequantMatricesRoundtrip(memory_manager, matrices)); 221 return true; 222 } 223 224 } // namespace jxl