enc_ma.h (6406B)
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 #ifndef LIB_JXL_MODULAR_ENCODING_ENC_MA_H_ 7 #define LIB_JXL_MODULAR_ENCODING_ENC_MA_H_ 8 9 #include <numeric> 10 11 #include "lib/jxl/enc_ans.h" 12 #include "lib/jxl/entropy_coder.h" 13 #include "lib/jxl/modular/encoding/dec_ma.h" 14 #include "lib/jxl/modular/modular_image.h" 15 #include "lib/jxl/modular/options.h" 16 17 namespace jxl { 18 19 // Struct to collect all the data needed to build a tree. 20 struct TreeSamples { 21 bool HasSamples() const { 22 return !residuals.empty() && !residuals[0].empty(); 23 } 24 size_t NumDistinctSamples() const { return sample_counts.size(); } 25 size_t NumSamples() const { return num_samples; } 26 // Set the predictor to use. Must be called before adding any samples. 27 Status SetPredictor(Predictor predictor, 28 ModularOptions::TreeMode wp_tree_mode); 29 // Set the properties to use. Must be called before adding any samples. 30 Status SetProperties(const std::vector<uint32_t> &properties, 31 ModularOptions::TreeMode wp_tree_mode); 32 33 size_t Token(size_t pred, size_t i) const { return residuals[pred][i].tok; } 34 size_t NBits(size_t pred, size_t i) const { return residuals[pred][i].nbits; } 35 size_t Count(size_t i) const { return sample_counts[i]; } 36 size_t PredictorIndex(Predictor predictor) const { 37 const auto predictor_elem = 38 std::find(predictors.begin(), predictors.end(), predictor); 39 JXL_DASSERT(predictor_elem != predictors.end()); 40 return predictor_elem - predictors.begin(); 41 } 42 size_t PropertyIndex(size_t property) const { 43 const auto property_elem = 44 std::find(props_to_use.begin(), props_to_use.end(), property); 45 JXL_DASSERT(property_elem != props_to_use.end()); 46 return property_elem - props_to_use.begin(); 47 } 48 size_t NumPropertyValues(size_t property_index) const { 49 return compact_properties[property_index].size() + 1; 50 } 51 // Returns the *quantized* property value. 52 size_t Property(size_t property_index, size_t i) const { 53 return props[property_index][i]; 54 } 55 int UnquantizeProperty(size_t property_index, uint32_t quant) const { 56 JXL_ASSERT(quant < compact_properties[property_index].size()); 57 return compact_properties[property_index][quant]; 58 } 59 60 Predictor PredictorFromIndex(size_t index) const { 61 JXL_DASSERT(index < predictors.size()); 62 return predictors[index]; 63 } 64 size_t PropertyFromIndex(size_t index) const { 65 JXL_DASSERT(index < props_to_use.size()); 66 return props_to_use[index]; 67 } 68 size_t NumPredictors() const { return predictors.size(); } 69 size_t NumProperties() const { return props_to_use.size(); } 70 71 // Preallocate data for a given number of samples. MUST be called before 72 // adding any sample. 73 void PrepareForSamples(size_t num_samples); 74 // Add a sample. 75 void AddSample(pixel_type_w pixel, const Properties &properties, 76 const pixel_type_w *predictions); 77 // Pre-cluster property values. 78 void PreQuantizeProperties( 79 const StaticPropRange &range, 80 const std::vector<ModularMultiplierInfo> &multiplier_info, 81 const std::vector<uint32_t> &group_pixel_count, 82 const std::vector<uint32_t> &channel_pixel_count, 83 std::vector<pixel_type> &pixel_samples, 84 std::vector<pixel_type> &diff_samples, size_t max_property_values); 85 86 void AllSamplesDone() { dedup_table_ = std::vector<uint32_t>(); } 87 88 uint32_t QuantizeProperty(uint32_t prop, pixel_type v) const { 89 v = std::min(std::max(v, -kPropertyRange), kPropertyRange) + kPropertyRange; 90 return property_mapping[prop][v]; 91 } 92 93 // Swaps samples in position a and b. Does nothing if a == b. 94 void Swap(size_t a, size_t b); 95 96 // Cycles samples: a -> b -> c -> a. We assume a <= b <= c, so that we can 97 // just call Swap(a, b) if b==c. 98 void ThreeShuffle(size_t a, size_t b, size_t c); 99 100 private: 101 // TODO(veluca): as the total number of properties and predictors are known 102 // before adding any samples, it might be better to interleave predictors, 103 // properties and counts in a single vector to improve locality. 104 // A first attempt at doing this actually results in much slower encoding, 105 // possibly because of the more complex addressing. 106 struct ResidualToken { 107 uint8_t tok; 108 uint8_t nbits; 109 }; 110 // Residual information: token and number of extra bits, per predictor. 111 std::vector<std::vector<ResidualToken>> residuals; 112 // Number of occurrences of each sample. 113 std::vector<uint16_t> sample_counts; 114 // Property values, quantized to at most 256 distinct values. 115 std::vector<std::vector<uint8_t>> props; 116 // Decompactification info for `props`. 117 std::vector<std::vector<int32_t>> compact_properties; 118 // List of properties to use. 119 std::vector<uint32_t> props_to_use; 120 // List of predictors to use. 121 std::vector<Predictor> predictors; 122 // Mapping property value -> quantized property value. 123 static constexpr int32_t kPropertyRange = 511; 124 std::vector<std::vector<uint8_t>> property_mapping; 125 // Number of samples seen. 126 size_t num_samples = 0; 127 // Table for deduplication. 128 static constexpr uint32_t kDedupEntryUnused{static_cast<uint32_t>(-1)}; 129 std::vector<uint32_t> dedup_table_; 130 131 // Functions for sample deduplication. 132 bool IsSameSample(size_t a, size_t b) const; 133 size_t Hash1(size_t a) const; 134 size_t Hash2(size_t a) const; 135 void InitTable(size_t size); 136 // Returns true if `a` was already present in the table. 137 bool AddToTableAndMerge(size_t a); 138 void AddToTable(size_t a); 139 }; 140 141 void TokenizeTree(const Tree &tree, std::vector<Token> *tokens, 142 Tree *decoder_tree); 143 144 void CollectPixelSamples(const Image &image, const ModularOptions &options, 145 size_t group_id, 146 std::vector<uint32_t> &group_pixel_count, 147 std::vector<uint32_t> &channel_pixel_count, 148 std::vector<pixel_type> &pixel_samples, 149 std::vector<pixel_type> &diff_samples); 150 151 void ComputeBestTree(TreeSamples &tree_samples, float threshold, 152 const std::vector<ModularMultiplierInfo> &mul_info, 153 StaticPropRange static_prop_range, 154 float fast_decode_multiplier, Tree *tree); 155 156 } // namespace jxl 157 #endif // LIB_JXL_MODULAR_ENCODING_ENC_MA_H_