third_party/jpeg-xl/lib/jpegli/color_quantize.cc

   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/jpegli/color_quantize.h"
   7
   8 #include <cmath>
   9 #include <limits>
  10 #include <unordered_map>
  11
  12 #include "lib/jpegli/decode_internal.h"
  13 #include "lib/jpegli/error.h"
  14
  15 namespace jpegli {
  16
  17 namespace {
  18
  19 static constexpr int kNumColorCellBits[kMaxComponents] = {3, 4, 3, 3};
  20 static constexpr int kCompW[kMaxComponents] = {2, 3, 1, 1};
  21
  22 int Pow(int a, int b) {
  23   int r = 1;
  24   for (int i = 0; i < b; ++i) {
  25     r *= a;
  26   }
  27   return r;
  28 }
  29
  30 int ComponentOrder(j_decompress_ptr cinfo, int i) {
  31   if (cinfo->out_color_components == 3) {
  32     return i < 2 ? 1 - i : i;
  33   }
  34   return i;
  35 }
  36
  37 int GetColorComponent(int i, int N) {
  38   return (i * 255 + (N - 1) / 2) / (N - 1);
  39 }
  40
  41 }  // namespace
  42
  43 void ChooseColorMap1Pass(j_decompress_ptr cinfo) {
  44   jpeg_decomp_master* m = cinfo->master;
  45   int components = cinfo->out_color_components;
  46   int desired = std::min(cinfo->desired_number_of_colors, 256);
  47   int num = 1;
  48   while (Pow(num + 1, components) <= desired) {
  49     ++num;
  50   }
  51   if (num == 1) {
  52     JPEGLI_ERROR("Too few colors (%d) in requested colormap", desired);
  53   }
  54   int actual = Pow(num, components);
  55   for (int i = 0; i < components; ++i) {
  56     m->num_colors_[i] = num;
  57   }
  58   while (actual < desired) {
  59     int total = actual;
  60     for (int i = 0; i < components; ++i) {
  61       int c = ComponentOrder(cinfo, i);
  62       int new_total = (actual / m->num_colors_[c]) * (m->num_colors_[c] + 1);
  63       if (new_total <= desired) {
  64         ++m->num_colors_[c];
  65         actual = new_total;
  66       }
  67     }
  68     if (actual == total) {
  69       break;
  70     }
  71   }
  72   cinfo->actual_number_of_colors = actual;
  73   cinfo->colormap = (*cinfo->mem->alloc_sarray)(
  74       reinterpret_cast<j_common_ptr>(cinfo), JPOOL_IMAGE, actual, components);
  75   int next_color[kMaxComponents] = {0};
  76   for (int i = 0; i < actual; ++i) {
  77     for (int c = 0; c < components; ++c) {
  78       cinfo->colormap[c][i] =
  79           GetColorComponent(next_color[c], m->num_colors_[c]);
  80     }
  81     int c = components - 1;
  82     while (c > 0 && next_color[c] + 1 == m->num_colors_[c]) {
  83       next_color[c--] = 0;
  84     }
  85     ++next_color[c];
  86   }
  87   if (!m->colormap_lut_) {
  88     m->colormap_lut_ = Allocate<uint8_t>(cinfo, components * 256, JPOOL_IMAGE);
  89   }
  90   int stride = actual;
  91   for (int c = 0; c < components; ++c) {
  92     int N = m->num_colors_[c];
  93     stride /= N;
  94     for (int i = 0; i < 256; ++i) {
  95       int index = ((2 * i - 1) * (N - 1) + 254) / 510;
  96       m->colormap_lut_[c * 256 + i] = index * stride;
  97     }
  98   }
  99 }
 100
 101 namespace {
 102
 103 // 2^13 priority levels for the PQ seems to be a good compromise between
 104 // accuracy, running time and stack space usage.
 105 static const int kMaxPriority = 1 << 13;
 106 static const int kMaxLevel = 3;
 107
 108 // This function is used in the multi-resolution grid to be able to compute
 109 // the keys for the different resolutions by just shifting the first key.
 110 inline int InterlaceBitsRGB(uint8_t r, uint8_t g, uint8_t b) {
 111   int z = 0;
 112   for (int i = 0; i < 7; ++i) {
 113     z += (r >> 5) & 4;
 114     z += (g >> 6) & 2;
 115     z += (b >> 7);
 116     z <<= 3;
 117     r <<= 1;
 118     g <<= 1;
 119     b <<= 1;
 120   }
 121   z += (r >> 5) & 4;
 122   z += (g >> 6) & 2;
 123   z += (b >> 7);
 124   return z;
 125 }
 126
 127 // This function will compute the actual priorities of the colors based on
 128 // the current distance from the palette, the population count and the signals
 129 // from the multi-resolution grid.
 130 inline int Priority(int d, int n, const int* density, const int* radius) {
 131   int p = d * n;
 132   for (int level = 0; level < kMaxLevel; ++level) {
 133     if (d > radius[level]) {
 134       p += density[level] * (d - radius[level]);
 135     }
 136   }
 137   return std::min(kMaxPriority - 1, p >> 4);
 138 }
 139
 140 inline int ColorIntQuadDistanceRGB(uint8_t r1, uint8_t g1, uint8_t b1,
 141                                    uint8_t r2, uint8_t g2, uint8_t b2) {
 142   // weights for the intensity calculation
 143   static constexpr int ired = 2;
 144   static constexpr int igreen = 5;
 145   static constexpr int iblue = 1;
 146   // normalization factor for the intensity calculation (2^ishift)
 147   static constexpr int ishift = 3;
 148   const int rd = r1 - r2;
 149   const int gd = g1 - g2;
 150   const int bd = b1 - b2;
 151   const int id = ired * rd + igreen * gd + iblue * bd;
 152   return rd * rd + gd * gd + bd * bd + ((id * id) >> (2 * ishift));
 153 }
 154
 155 inline int ScaleQuadDistanceRGB(int d) {
 156   return static_cast<int>(sqrt(d * 0.25) + 0.5);
 157 }
 158
 159 // The function updates the minimal distances, the clustering and the
 160 // quantization error after the insertion of the new color into the palette.
 161 void AddToRGBPalette(const uint8_t* red, const uint8_t* green,
 162                      const uint8_t* blue,
 163                      const int* count,  // histogram of colors
 164                      const int index,   // index of color to be added
 165                      const int k,       // size of current palette
 166                      const int n,       // number of colors
 167                      int* dist,         // array of distances from palette
 168                      int* cluster,      // mapping of color indices to palette
 169                      int* center,       // the inverse mapping
 170                      int64_t* error) {  // measure of the quantization error
 171   center[k] = index;
 172   cluster[index] = k;
 173   *error -=
 174       static_cast<int64_t>(dist[index]) * static_cast<int64_t>(count[index]);
 175   dist[index] = 0;
 176   for (int j = 0; j < n; ++j) {
 177     if (dist[j] > 0) {
 178       const int d = ColorIntQuadDistanceRGB(
 179           red[index], green[index], blue[index], red[j], green[j], blue[j]);
 180       if (d < dist[j]) {
 181         *error += static_cast<int64_t>((d - dist[j])) *
 182                   static_cast<int64_t>(count[j]);
 183         dist[j] = d;
 184         cluster[j] = k;
 185       }
 186     }
 187   }
 188 }
 189
 190 struct RGBPixelHasher {
 191   // A quick but good-enough hash to get 24 bits of RGB into the lower 12 bits.
 192   size_t operator()(uint32_t a) const { return (a ^ (a >> 12)) * 0x9e3779b9; }
 193 };
 194
 195 struct WangHasher {
 196   // Thomas Wang's Hash.  Nearly perfect and still quite fast.  Above (for
 197   // pixels) we use a simpler hash because the number of hash calls is
 198   // proportional to the number of pixels and that hash dominates; we want the
 199   // cost to be minimal and we start with a large table.  We can use a better
 200   // hash for the histogram since the number of hash calls is proportional to
 201   // the number of unique colors in the image, which is hopefully much smaller.
 202   // Note that the difference is slight; e.g. replacing RGBPixelHasher with
 203   // WangHasher only slows things down by 5% on an Opteron.
 204   size_t operator()(uint32_t a) const {
 205     a = (a ^ 61) ^ (a >> 16);
 206     a = a + (a << 3);
 207     a = a ^ (a >> 4);
 208     a = a * 0x27d4eb2d;
 209     a = a ^ (a >> 15);
 210     return a;
 211   }
 212 };
 213
 214 // Build an index of all the different colors in the input
 215 // image. To do this we map the 24 bit RGB representation of the colors
 216 // to a unique integer index assigned to the different colors in order of
 217 // appearance in the image.  Return the number of unique colors found.
 218 // The colors are pre-quantized to 3 * 6 bits precision.
 219 static int BuildRGBColorIndex(const uint8_t* const image, int const num_pixels,
 220                               int* const count, uint8_t* const red,
 221                               uint8_t* const green, uint8_t* const blue) {
 222   // Impossible because rgb are in the low 24 bits, and the upper 8 bits is 0.
 223   const uint32_t impossible_pixel_value = 0x10000000;
 224   std::unordered_map<uint32_t, int, RGBPixelHasher> index_map(1 << 12);
 225   std::unordered_map<uint32_t, int, RGBPixelHasher>::iterator index_map_lookup;
 226   const uint8_t* imagep = &image[0];
 227   uint32_t prev_pixel = impossible_pixel_value;
 228   int index = 0;
 229   int n = 0;
 230   for (int i = 0; i < num_pixels; ++i) {
 231     uint8_t r = ((*imagep++) & 0xfc) + 2;
 232     uint8_t g = ((*imagep++) & 0xfc) + 2;
 233     uint8_t b = ((*imagep++) & 0xfc) + 2;
 234     uint32_t pixel = (b << 16) | (g << 8) | r;
 235     if (pixel != prev_pixel) {
 236       prev_pixel = pixel;
 237       index_map_lookup = index_map.find(pixel);
 238       if (index_map_lookup != index_map.end()) {
 239         index = index_map_lookup->second;
 240       } else {
 241         index_map[pixel] = index = n++;
 242         red[index] = r;
 243         green[index] = g;
 244         blue[index] = b;
 245       }
 246     }
 247     ++count[index];
 248   }
 249   return n;
 250 }
 251
 252 }  // namespace
 253
 254 void ChooseColorMap2Pass(j_decompress_ptr cinfo) {
 255   if (cinfo->out_color_space != JCS_RGB) {
 256     JPEGLI_ERROR("Two-pass quantizer must use RGB output color space.");
 257   }
 258   jpeg_decomp_master* m = cinfo->master;
 259   const size_t num_pixels = cinfo->output_width * cinfo->output_height;
 260   const int max_color_count = std::max<size_t>(num_pixels, 1u << 18);
 261   const int max_palette_size = cinfo->desired_number_of_colors;
 262   std::unique_ptr<uint8_t[]> red(new uint8_t[max_color_count]);
 263   std::unique_ptr<uint8_t[]> green(new uint8_t[max_color_count]);
 264   std::unique_ptr<uint8_t[]> blue(new uint8_t[max_color_count]);
 265   std::vector<int> count(max_color_count, 0);
 266   // number of colors
 267   int n = BuildRGBColorIndex(m->pixels_, num_pixels, &count[0], &red[0],
 268                              &green[0], &blue[0]);
 269
 270   std::vector<int> dist(n, std::numeric_limits<int>::max());
 271   std::vector<int> cluster(n);
 272   std::vector<bool> in_palette(n, false);
 273   int center[256];
 274   int k = 0;  // palette size
 275   const int count_threshold = (num_pixels * 4) / max_palette_size;
 276   static constexpr int kAveragePixelErrorThreshold = 1;
 277   const int64_t error_threshold = num_pixels * kAveragePixelErrorThreshold;
 278   int64_t error = 0;  // quantization error
 279
 280   int max_count = 0;
 281   int winner = 0;
 282   for (int i = 0; i < n; ++i) {
 283     if (count[i] > max_count) {
 284       max_count = count[i];
 285       winner = i;
 286     }
 287     if (!in_palette[i] && count[i] > count_threshold) {
 288       AddToRGBPalette(&red[0], &green[0], &blue[0], &count[0], i, k++, n,
 289                       &dist[0], &cluster[0], &center[0], &error);
 290       in_palette[i] = true;
 291     }
 292   }
 293   if (k == 0) {
 294     AddToRGBPalette(&red[0], &green[0], &blue[0], &count[0], winner, k++, n,
 295                     &dist[0], &cluster[0], &center[0], &error);
 296     in_palette[winner] = true;
 297   }
 298
 299   // Calculation of the multi-resolution density grid.
 300   std::vector<int> density(n * kMaxLevel);
 301   std::vector<int> radius(n * kMaxLevel);
 302   std::unordered_map<uint32_t, int, WangHasher> histogram[kMaxLevel];
 303   for (int level = 0; level < kMaxLevel; ++level) {
 304     // This value is never used because key = InterlaceBitsRGB(...) >> 6
 305   }
 306
 307   for (int i = 0; i < n; ++i) {
 308     if (!in_palette[i]) {
 309       const int key = InterlaceBitsRGB(red[i], green[i], blue[i]) >> 6;
 310       for (int level = 0; level < kMaxLevel; ++level) {
 311         histogram[level][key >> (3 * level)] += count[i];
 312       }
 313     }
 314   }
 315   for (int i = 0; i < n; ++i) {
 316     if (!in_palette[i]) {
 317       for (int level = 0; level < kMaxLevel; ++level) {
 318         const int mask = (4 << level) - 1;
 319         const int rd = std::max(red[i] & mask, mask - (red[i] & mask));
 320         const int gd = std::max(green[i] & mask, mask - (green[i] & mask));
 321         const int bd = std::max(blue[i] & mask, mask - (blue[i] & mask));
 322         radius[i * kMaxLevel + level] =
 323             ScaleQuadDistanceRGB(ColorIntQuadDistanceRGB(0, 0, 0, rd, gd, bd));
 324       }
 325       const int key = InterlaceBitsRGB(red[i], green[i], blue[i]) >> 6;
 326       if (kMaxLevel > 0) {
 327         density[i * kMaxLevel] = histogram[0][key] - count[i];
 328       }
 329       for (int level = 1; level < kMaxLevel; ++level) {
 330         density[i * kMaxLevel + level] =
 331             (histogram[level][key >> (3 * level)] -
 332              histogram[level - 1][key >> (3 * level - 3)]);
 333       }
 334     }
 335   }
 336
 337   // Calculate the initial error now that the palette has been initialized.
 338   error = 0;
 339   for (int i = 0; i < n; ++i) {
 340     error += static_cast<int64_t>(dist[i]) * static_cast<int64_t>(count[i]);
 341   }
 342
 343   std::unique_ptr<std::vector<int>[]> bucket_array(
 344       new std::vector<int>[kMaxPriority]);
 345   int top_priority = -1;
 346   for (int i = 0; i < n; ++i) {
 347     if (!in_palette[i]) {
 348       int priority = Priority(ScaleQuadDistanceRGB(dist[i]), count[i],
 349                               &density[i * kMaxLevel], &radius[i * kMaxLevel]);
 350       bucket_array[priority].push_back(i);
 351       top_priority = std::max(priority, top_priority);
 352     }
 353   }
 354   double error_accum = 0;
 355   while (top_priority >= 0 && k < max_palette_size) {
 356     if (error < error_threshold) {
 357       error_accum += std::min(error_threshold, error_threshold - error);
 358       if (error_accum >= 10 * error_threshold) {
 359         break;
 360       }
 361     }
 362     int i = bucket_array[top_priority].back();
 363     int priority = Priority(ScaleQuadDistanceRGB(dist[i]), count[i],
 364                             &density[i * kMaxLevel], &radius[i * kMaxLevel]);
 365     if (priority < top_priority) {
 366       bucket_array[priority].push_back(i);
 367     } else {
 368       AddToRGBPalette(&red[0], &green[0], &blue[0], &count[0], i, k++, n,
 369                       &dist[0], &cluster[0], &center[0], &error);
 370     }
 371     bucket_array[top_priority].pop_back();
 372     while (top_priority >= 0 && bucket_array[top_priority].empty()) {
 373       --top_priority;
 374     }
 375   }
 376
 377   cinfo->actual_number_of_colors = k;
 378   cinfo->colormap = (*cinfo->mem->alloc_sarray)(
 379       reinterpret_cast<j_common_ptr>(cinfo), JPOOL_IMAGE, k, 3);
 380   for (int i = 0; i < k; ++i) {
 381     int index = center[i];
 382     cinfo->colormap[0][i] = red[index];
 383     cinfo->colormap[1][i] = green[index];
 384     cinfo->colormap[2][i] = blue[index];
 385   }
 386 }
 387
 388 namespace {
 389
 390 void FindCandidatesForCell(j_decompress_ptr cinfo, int ncomp, int cell[],
 391                            std::vector<uint8_t>* candidates) {
 392   int cell_min[kMaxComponents];
 393   int cell_max[kMaxComponents];
 394   int cell_center[kMaxComponents];
 395   for (int c = 0; c < ncomp; ++c) {
 396     cell_min[c] = cell[c] << (8 - kNumColorCellBits[c]);
 397     cell_max[c] = cell_min[c] + (1 << (8 - kNumColorCellBits[c])) - 1;
 398     cell_center[c] = (cell_min[c] + cell_max[c]) >> 1;
 399   }
 400   int min_maxdist = std::numeric_limits<int>::max();
 401   int mindist[256];
 402   for (int i = 0; i < cinfo->actual_number_of_colors; ++i) {
 403     int dmin = 0;
 404     int dmax = 0;
 405     for (int c = 0; c < ncomp; ++c) {
 406       int palette_c = cinfo->colormap[c][i];
 407       int dminc = 0, dmaxc;
 408       if (palette_c < cell_min[c]) {
 409         dminc = cell_min[c] - palette_c;
 410         dmaxc = cell_max[c] - palette_c;
 411       } else if (palette_c > cell_max[c]) {
 412         dminc = palette_c - cell_max[c];
 413         dmaxc = palette_c - cell_min[c];
 414       } else if (palette_c > cell_center[c]) {
 415         dmaxc = palette_c - cell_min[c];
 416       } else {
 417         dmaxc = cell_max[c] - palette_c;
 418       }
 419       dminc *= kCompW[c];
 420       dmaxc *= kCompW[c];
 421       dmin += dminc * dminc;
 422       dmax += dmaxc * dmaxc;
 423     }
 424     mindist[i] = dmin;
 425     min_maxdist = std::min(dmax, min_maxdist);
 426   }
 427   for (int i = 0; i < cinfo->actual_number_of_colors; ++i) {
 428     if (mindist[i] < min_maxdist) {
 429       candidates->push_back(i);
 430     }
 431   }
 432 }
 433
 434 }  // namespace
 435
 436 void CreateInverseColorMap(j_decompress_ptr cinfo) {
 437   jpeg_decomp_master* m = cinfo->master;
 438   int ncomp = cinfo->out_color_components;
 439   int num_cells = 1;
 440   for (int c = 0; c < ncomp; ++c) {
 441     num_cells *= (1 << kNumColorCellBits[c]);
 442   }
 443   m->candidate_lists_.resize(num_cells);
 444
 445   int next_cell[kMaxComponents] = {0};
 446   for (int i = 0; i < num_cells; ++i) {
 447     m->candidate_lists_[i].clear();
 448     FindCandidatesForCell(cinfo, ncomp, next_cell, &m->candidate_lists_[i]);
 449     int c = ncomp - 1;
 450     while (c > 0 && next_cell[c] + 1 == (1 << kNumColorCellBits[c])) {
 451       next_cell[c--] = 0;
 452     }
 453     ++next_cell[c];
 454   }
 455   m->regenerate_inverse_colormap_ = false;
 456 }
 457
 458 int LookupColorIndex(j_decompress_ptr cinfo, JSAMPLE* pixel) {
 459   jpeg_decomp_master* m = cinfo->master;
 460   int num_channels = cinfo->out_color_components;
 461   int index = 0;
 462   if (m->quant_mode_ == 1) {
 463     for (int c = 0; c < num_channels; ++c) {
 464       index += m->colormap_lut_[c * 256 + pixel[c]];
 465     }
 466   } else {
 467     size_t cell_idx = 0;
 468     size_t stride = 1;
 469     for (int c = num_channels - 1; c >= 0; --c) {
 470       cell_idx += (pixel[c] >> (8 - kNumColorCellBits[c])) * stride;
 471       stride <<= kNumColorCellBits[c];
 472     }
 473     JXL_ASSERT(cell_idx < m->candidate_lists_.size());
 474     int mindist = std::numeric_limits<int>::max();
 475     const auto& candidates = m->candidate_lists_[cell_idx];
 476     for (uint8_t i : candidates) {
 477       int dist = 0;
 478       for (int c = 0; c < num_channels; ++c) {
 479         int d = (cinfo->colormap[c][i] - pixel[c]) * kCompW[c];
 480         dist += d * d;
 481       }
 482       if (dist < mindist) {
 483         mindist = dist;
 484         index = i;
 485       }
 486     }
 487   }
 488   JXL_ASSERT(index < cinfo->actual_number_of_colors);
 489   return index;
 490 }
 491
 492 void CreateOrderedDitherTables(j_decompress_ptr cinfo) {
 493   jpeg_decomp_master* m = cinfo->master;
 494   static constexpr size_t kDitherSize = 4;
 495   static constexpr size_t kDitherMask = kDitherSize - 1;
 496   static constexpr float kBaseDitherMatrix[] = {
 497       0,  8,  2,  10,  //
 498       12, 4,  14, 6,   //
 499       3,  11, 1,  9,   //
 500       15, 7,  13, 5,   //
 501   };
 502   m->dither_size_ = kDitherSize;
 503   m->dither_mask_ = kDitherMask;
 504   size_t ncells = m->dither_size_ * m->dither_size_;
 505   for (int c = 0; c < cinfo->out_color_components; ++c) {
 506     float spread = 1.0f / (m->num_colors_[c] - 1);
 507     float mul = spread / ncells;
 508     float offset = 0.5f * spread;
 509     if (m->dither_[c] == nullptr) {
 510       m->dither_[c] = Allocate<float>(cinfo, ncells, JPOOL_IMAGE_ALIGNED);
 511     }
 512     for (size_t idx = 0; idx < ncells; ++idx) {
 513       m->dither_[c][idx] = kBaseDitherMatrix[idx] * mul - offset;
 514     }
 515   }
 516 }
 517
 518 void InitFSDitherState(j_decompress_ptr cinfo) {
 519   jpeg_decomp_master* m = cinfo->master;
 520   for (int c = 0; c < cinfo->out_color_components; ++c) {
 521     if (m->error_row_[c] == nullptr) {
 522       m->error_row_[c] =
 523           Allocate<float>(cinfo, cinfo->output_width, JPOOL_IMAGE_ALIGNED);
 524       m->error_row_[c + kMaxComponents] =
 525           Allocate<float>(cinfo, cinfo->output_width, JPOOL_IMAGE_ALIGNED);
 526     }
 527     memset(m->error_row_[c], 0.0, cinfo->output_width * sizeof(float));
 528     memset(m->error_row_[c + kMaxComponents], 0.0,
 529            cinfo->output_width * sizeof(float));
 530   }
 531 }
 532
 533 }  // namespace jpegli