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