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/jpegli/color_quantize.h"
#include <cmath>
#include <limits>
#include <unordered_map>
#include "lib/jpegli/decode_internal.h"
#include "lib/jpegli/error.h"
namespace jpegli {
namespace {
constexpr int kNumColorCellBits[kMaxComponents] = {3, 4, 3, 3};
constexpr int kCompW[kMaxComponents] = {2, 3, 1, 1};
int Pow(int a, int b) {
int r = 1;
for (int i = 0; i < b; ++i) {
r *= a;
}
return r;
}
int ComponentOrder(j_decompress_ptr cinfo, int i) {
if (cinfo->out_color_components == 3) {
return i < 2 ? 1 - i : i;
}
return i;
}
int GetColorComponent(int i, int N) {
return (i * 255 + (N - 1) / 2) / (N - 1);
}
} // namespace
void ChooseColorMap1Pass(j_decompress_ptr cinfo) {
jpeg_decomp_master* m = cinfo->master;
int components = cinfo->out_color_components;
int desired = std::min(cinfo->desired_number_of_colors, 256);
int num = 1;
while (Pow(num + 1, components) <= desired) {
++num;
}
if (num == 1) {
JPEGLI_ERROR("Too few colors (%d) in requested colormap", desired);
}
int actual = Pow(num, components);
for (int i = 0; i < components; ++i) {
m->num_colors_[i] = num;
}
while (actual < desired) {
int total = actual;
for (int i = 0; i < components; ++i) {
int c = ComponentOrder(cinfo, i);
int new_total = (actual / m->num_colors_[c]) * (m->num_colors_[c] + 1);
if (new_total <= desired) {
++m->num_colors_[c];
actual = new_total;
}
}
if (actual == total) {
break;
}
}
cinfo->actual_number_of_colors = actual;
cinfo->colormap = (*cinfo->mem->alloc_sarray)(
reinterpret_cast<j_common_ptr>(cinfo), JPOOL_IMAGE, actual, components);
int next_color[kMaxComponents] = {0};
for (int i = 0; i < actual; ++i) {
for (int c = 0; c < components; ++c) {
cinfo->colormap[c][i] =
GetColorComponent(next_color[c], m->num_colors_[c]);
}
int c = components - 1;
while (c > 0 && next_color[c] + 1 == m->num_colors_[c]) {
next_color[c--] = 0;
}
++next_color[c];
}
if (!m->colormap_lut_) {
m->colormap_lut_ = Allocate<uint8_t>(cinfo, components * 256, JPOOL_IMAGE);
}
int stride = actual;
for (int c = 0; c < components; ++c) {
int N = m->num_colors_[c];
stride /= N;
for (int i = 0; i < 256; ++i) {
int index = ((2 * i - 1) * (N - 1) + 254) / 510;
m->colormap_lut_[c * 256 + i] = index * stride;
}
}
}
namespace {
// 2^13 priority levels for the PQ seems to be a good compromise between
// accuracy, running time and stack space usage.
const int kMaxPriority = 1 << 13;
const int kMaxLevel = 3;
// This function is used in the multi-resolution grid to be able to compute
// the keys for the different resolutions by just shifting the first key.
inline int InterlaceBitsRGB(uint8_t r, uint8_t g, uint8_t b) {
int z = 0;
for (int i = 0; i < 7; ++i) {
z += (r >> 5) & 4;
z += (g >> 6) & 2;
z += (b >> 7);
z <<= 3;
r <<= 1;
g <<= 1;
b <<= 1;
}
z += (r >> 5) & 4;
z += (g >> 6) & 2;
z += (b >> 7);
return z;
}
// This function will compute the actual priorities of the colors based on
// the current distance from the palette, the population count and the signals
// from the multi-resolution grid.
inline int Priority(int d, int n, const int* density, const int* radius) {
int p = d * n;
for (int level = 0; level < kMaxLevel; ++level) {
if (d > radius[level]) {
p += density[level] * (d - radius[level]);
}
}
return std::min(kMaxPriority - 1, p >> 4);
}
inline int ColorIntQuadDistanceRGB(uint8_t r1, uint8_t g1, uint8_t b1,
uint8_t r2, uint8_t g2, uint8_t b2) {
// weights for the intensity calculation
static constexpr int ired = 2;
static constexpr int igreen = 5;
static constexpr int iblue = 1;
// normalization factor for the intensity calculation (2^ishift)
static constexpr int ishift = 3;
const int rd = r1 - r2;
const int gd = g1 - g2;
const int bd = b1 - b2;
const int id = ired * rd + igreen * gd + iblue * bd;
return rd * rd + gd * gd + bd * bd + ((id * id) >> (2 * ishift));
}
inline int ScaleQuadDistanceRGB(int d) {
return static_cast<int>(std::lround(sqrt(d * 0.25)));
}
// The function updates the minimal distances, the clustering and the
// quantization error after the insertion of the new color into the palette.
void AddToRGBPalette(const uint8_t* red, const uint8_t* green,
const uint8_t* blue,
const int* count, // histogram of colors
const int index, // index of color to be added
const int k, // size of current palette
const int n, // number of colors
int* dist, // array of distances from palette
int* cluster, // mapping of color indices to palette
int* center, // the inverse mapping
int64_t* error) { // measure of the quantization error
center[k] = index;
cluster[index] = k;
*error -=
static_cast<int64_t>(dist[index]) * static_cast<int64_t>(count[index]);
dist[index] = 0;
for (int j = 0; j < n; ++j) {
if (dist[j] > 0) {
const int d = ColorIntQuadDistanceRGB(
red[index], green[index], blue[index], red[j], green[j], blue[j]);
if (d < dist[j]) {
*error += static_cast<int64_t>((d - dist[j])) *
static_cast<int64_t>(count[j]);
dist[j] = d;
cluster[j] = k;
}
}
}
}
struct RGBPixelHasher {
// A quick but good-enough hash to get 24 bits of RGB into the lower 12 bits.
size_t operator()(uint32_t a) const { return (a ^ (a >> 12)) * 0x9e3779b9; }
};
struct WangHasher {
// Thomas Wang's Hash. Nearly perfect and still quite fast. Above (for
// pixels) we use a simpler hash because the number of hash calls is
// proportional to the number of pixels and that hash dominates; we want the
// cost to be minimal and we start with a large table. We can use a better
// hash for the histogram since the number of hash calls is proportional to
// the number of unique colors in the image, which is hopefully much smaller.
// Note that the difference is slight; e.g. replacing RGBPixelHasher with
// WangHasher only slows things down by 5% on an Opteron.
size_t operator()(uint32_t a) const {
a = (a ^ 61) ^ (a >> 16);
a = a + (a << 3);
a = a ^ (a >> 4);
a = a * 0x27d4eb2d;
a = a ^ (a >> 15);
return a;
}
};
// Build an index of all the different colors in the input
// image. To do this we map the 24 bit RGB representation of the colors
// to a unique integer index assigned to the different colors in order of
// appearance in the image. Return the number of unique colors found.
// The colors are pre-quantized to 3 * 6 bits precision.
int BuildRGBColorIndex(const uint8_t* const image, int const num_pixels,
int* const count, uint8_t* const red,
uint8_t* const green, uint8_t* const blue) {
// Impossible because rgb are in the low 24 bits, and the upper 8 bits is 0.
const uint32_t impossible_pixel_value = 0x10000000;
std::unordered_map<uint32_t, int, RGBPixelHasher> index_map(1 << 12);
std::unordered_map<uint32_t, int, RGBPixelHasher>::iterator index_map_lookup;
const uint8_t* imagep = &image[0];
uint32_t prev_pixel = impossible_pixel_value;
int index = 0;
int n = 0;
for (int i = 0; i < num_pixels; ++i) {
uint8_t r = ((*imagep++) & 0xfc) + 2;
uint8_t g = ((*imagep++) & 0xfc) + 2;
uint8_t b = ((*imagep++) & 0xfc) + 2;
uint32_t pixel = (b << 16) | (g << 8) | r;
if (pixel != prev_pixel) {
prev_pixel = pixel;
index_map_lookup = index_map.find(pixel);
if (index_map_lookup != index_map.end()) {
index = index_map_lookup->second;
} else {
index_map[pixel] = index = n++;
red[index] = r;
green[index] = g;
blue[index] = b;
}
}
++count[index];
}
return n;
}
} // namespace
void ChooseColorMap2Pass(j_decompress_ptr cinfo) {
if (cinfo->out_color_space != JCS_RGB) {
JPEGLI_ERROR("Two-pass quantizer must use RGB output color space.");
}
jpeg_decomp_master* m = cinfo->master;
const size_t num_pixels = cinfo->output_width * cinfo->output_height;
const int max_color_count = std::max<size_t>(num_pixels, 1u << 18);
const int max_palette_size = cinfo->desired_number_of_colors;
std::unique_ptr<uint8_t[]> red(new uint8_t[max_color_count]);
std::unique_ptr<uint8_t[]> green(new uint8_t[max_color_count]);
std::unique_ptr<uint8_t[]> blue(new uint8_t[max_color_count]);
std::vector<int> count(max_color_count, 0);
// number of colors
int n = BuildRGBColorIndex(m->pixels_, num_pixels, count.data(), &red[0],
&green[0], &blue[0]);
std::vector<int> dist(n, std::numeric_limits<int>::max());
std::vector<int> cluster(n);
std::vector<bool> in_palette(n, false);
int center[256];
int k = 0; // palette size
const int count_threshold = (num_pixels * 4) / max_palette_size;
static constexpr int kAveragePixelErrorThreshold = 1;
const int64_t error_threshold = num_pixels * kAveragePixelErrorThreshold;
int64_t error = 0; // quantization error
int max_count = 0;
int winner = 0;
for (int i = 0; i < n; ++i) {
if (count[i] > max_count) {
max_count = count[i];
winner = i;
}
if (!in_palette[i] && count[i] > count_threshold) {
AddToRGBPalette(&red[0], &green[0], &blue[0], count.data(), i, k++, n,
dist.data(), cluster.data(), ¢er[0], &error);
in_palette[i] = true;
}
}
if (k == 0) {
AddToRGBPalette(&red[0], &green[0], &blue[0], count.data(), winner, k++, n,
dist.data(), cluster.data(), ¢er[0], &error);
in_palette[winner] = true;
}
// Calculation of the multi-resolution density grid.
std::vector<int> density(n * kMaxLevel);
std::vector<int> radius(n * kMaxLevel);
std::unordered_map<uint32_t, int, WangHasher> histogram[kMaxLevel];
for (int level = 0; level < kMaxLevel; ++level) {
// This value is never used because key = InterlaceBitsRGB(...) >> 6
}
for (int i = 0; i < n; ++i) {
if (!in_palette[i]) {
const int key = InterlaceBitsRGB(red[i], green[i], blue[i]) >> 6;
for (int level = 0; level < kMaxLevel; ++level) {
histogram[level][key >> (3 * level)] += count[i];
}
}
}
for (int i = 0; i < n; ++i) {
if (!in_palette[i]) {
for (int level = 0; level < kMaxLevel; ++level) {
const int mask = (4 << level) - 1;
const int rd = std::max(red[i] & mask, mask - (red[i] & mask));
const int gd = std::max(green[i] & mask, mask - (green[i] & mask));
const int bd = std::max(blue[i] & mask, mask - (blue[i] & mask));
radius[i * kMaxLevel + level] =
ScaleQuadDistanceRGB(ColorIntQuadDistanceRGB(0, 0, 0, rd, gd, bd));
}
const int key = InterlaceBitsRGB(red[i], green[i], blue[i]) >> 6;
if (kMaxLevel > 0) {
density[i * kMaxLevel] = histogram[0][key] - count[i];
}
for (int level = 1; level < kMaxLevel; ++level) {
density[i * kMaxLevel + level] =
(histogram[level][key >> (3 * level)] -
histogram[level - 1][key >> (3 * level - 3)]);
}
}
}
// Calculate the initial error now that the palette has been initialized.
error = 0;
for (int i = 0; i < n; ++i) {
error += static_cast<int64_t>(dist[i]) * static_cast<int64_t>(count[i]);
}
std::unique_ptr<std::vector<int>[]> bucket_array(
new std::vector<int>[kMaxPriority]);
int top_priority = -1;
for (int i = 0; i < n; ++i) {
if (!in_palette[i]) {
int priority = Priority(ScaleQuadDistanceRGB(dist[i]), count[i],
&density[i * kMaxLevel], &radius[i * kMaxLevel]);
bucket_array[priority].push_back(i);
top_priority = std::max(priority, top_priority);
}
}
double error_accum = 0;
while (top_priority >= 0 && k < max_palette_size) {
if (error < error_threshold) {
error_accum += std::min(error_threshold, error_threshold - error);
if (error_accum >= 10 * error_threshold) {
break;
}
}
int i = bucket_array[top_priority].back();
int priority = Priority(ScaleQuadDistanceRGB(dist[i]), count[i],
&density[i * kMaxLevel], &radius[i * kMaxLevel]);
if (priority < top_priority) {
bucket_array[priority].push_back(i);
} else {
AddToRGBPalette(&red[0], &green[0], &blue[0], count.data(), i, k++, n,
dist.data(), cluster.data(), ¢er[0], &error);
}
bucket_array[top_priority].pop_back();
while (top_priority >= 0 && bucket_array[top_priority].empty()) {
--top_priority;
}
}
cinfo->actual_number_of_colors = k;
cinfo->colormap = (*cinfo->mem->alloc_sarray)(
reinterpret_cast<j_common_ptr>(cinfo), JPOOL_IMAGE, k, 3);
for (int i = 0; i < k; ++i) {
int index = center[i];
cinfo->colormap[0][i] = red[index];
cinfo->colormap[1][i] = green[index];
cinfo->colormap[2][i] = blue[index];
}
}
namespace {
void FindCandidatesForCell(j_decompress_ptr cinfo, int ncomp, const int cell[],
std::vector<uint8_t>* candidates) {
int cell_min[kMaxComponents];
int cell_max[kMaxComponents];
int cell_center[kMaxComponents];
for (int c = 0; c < ncomp; ++c) {
cell_min[c] = cell[c] << (8 - kNumColorCellBits[c]);
cell_max[c] = cell_min[c] + (1 << (8 - kNumColorCellBits[c])) - 1;
cell_center[c] = (cell_min[c] + cell_max[c]) >> 1;
}
int min_maxdist = std::numeric_limits<int>::max();
int mindist[256];
for (int i = 0; i < cinfo->actual_number_of_colors; ++i) {
int dmin = 0;
int dmax = 0;
for (int c = 0; c < ncomp; ++c) {
int palette_c = cinfo->colormap[c][i];
int dminc = 0;
int dmaxc;
if (palette_c < cell_min[c]) {
dminc = cell_min[c] - palette_c;
dmaxc = cell_max[c] - palette_c;
} else if (palette_c > cell_max[c]) {
dminc = palette_c - cell_max[c];
dmaxc = palette_c - cell_min[c];
} else if (palette_c > cell_center[c]) {
dmaxc = palette_c - cell_min[c];
} else {
dmaxc = cell_max[c] - palette_c;
}
dminc *= kCompW[c];
dmaxc *= kCompW[c];
dmin += dminc * dminc;
dmax += dmaxc * dmaxc;
}
mindist[i] = dmin;
min_maxdist = std::min(dmax, min_maxdist);
}
for (int i = 0; i < cinfo->actual_number_of_colors; ++i) {
if (mindist[i] < min_maxdist) {
candidates->push_back(i);
}
}
}
} // namespace
void CreateInverseColorMap(j_decompress_ptr cinfo) {
jpeg_decomp_master* m = cinfo->master;
int ncomp = cinfo->out_color_components;
JPEGLI_CHECK(ncomp > 0);
JPEGLI_CHECK(ncomp <= kMaxComponents);
int num_cells = 1;
for (int c = 0; c < ncomp; ++c) {
num_cells *= (1 << kNumColorCellBits[c]);
}
m->candidate_lists_.resize(num_cells);
int next_cell[kMaxComponents] = {0};
for (int i = 0; i < num_cells; ++i) {
m->candidate_lists_[i].clear();
FindCandidatesForCell(cinfo, ncomp, next_cell, &m->candidate_lists_[i]);
int c = ncomp - 1;
while (c > 0 && next_cell[c] + 1 == (1 << kNumColorCellBits[c])) {
next_cell[c--] = 0;
}
++next_cell[c];
}
m->regenerate_inverse_colormap_ = false;
}
int LookupColorIndex(j_decompress_ptr cinfo, const JSAMPLE* pixel) {
jpeg_decomp_master* m = cinfo->master;
int num_channels = cinfo->out_color_components;
int index = 0;
if (m->quant_mode_ == 1) {
for (int c = 0; c < num_channels; ++c) {
index += m->colormap_lut_[c * 256 + pixel[c]];
}
} else {
size_t cell_idx = 0;
size_t stride = 1;
for (int c = num_channels - 1; c >= 0; --c) {
cell_idx += (pixel[c] >> (8 - kNumColorCellBits[c])) * stride;
stride <<= kNumColorCellBits[c];
}
JPEGLI_CHECK(cell_idx < m->candidate_lists_.size());
int mindist = std::numeric_limits<int>::max();
const auto& candidates = m->candidate_lists_[cell_idx];
for (uint8_t i : candidates) {
int dist = 0;
for (int c = 0; c < num_channels; ++c) {
int d = (cinfo->colormap[c][i] - pixel[c]) * kCompW[c];
dist += d * d;
}
if (dist < mindist) {
mindist = dist;
index = i;
}
}
}
JPEGLI_CHECK(index < cinfo->actual_number_of_colors);
return index;
}
void CreateOrderedDitherTables(j_decompress_ptr cinfo) {
jpeg_decomp_master* m = cinfo->master;
static constexpr size_t kDitherSize = 4;
static constexpr size_t kDitherMask = kDitherSize - 1;
static constexpr float kBaseDitherMatrix[] = {
0, 8, 2, 10, //
12, 4, 14, 6, //
3, 11, 1, 9, //
15, 7, 13, 5, //
};
m->dither_size_ = kDitherSize;
m->dither_mask_ = kDitherMask;
size_t ncells = m->dither_size_ * m->dither_size_;
for (int c = 0; c < cinfo->out_color_components; ++c) {
float spread = 1.0f / (m->num_colors_[c] - 1);
float mul = spread / ncells;
float offset = 0.5f * spread;
if (m->dither_[c] == nullptr) {
m->dither_[c] = Allocate<float>(cinfo, ncells, JPOOL_IMAGE_ALIGNED);
}
for (size_t idx = 0; idx < ncells; ++idx) {
m->dither_[c][idx] = kBaseDitherMatrix[idx] * mul - offset;
}
}
}
void InitFSDitherState(j_decompress_ptr cinfo) {
jpeg_decomp_master* m = cinfo->master;
for (int c = 0; c < cinfo->out_color_components; ++c) {
if (m->error_row_[c] == nullptr) {
m->error_row_[c] =
Allocate<float>(cinfo, cinfo->output_width, JPOOL_IMAGE_ALIGNED);
m->error_row_[c + kMaxComponents] =
Allocate<float>(cinfo, cinfo->output_width, JPOOL_IMAGE_ALIGNED);
}
memset(m->error_row_[c], 0.0, cinfo->output_width * sizeof(float));
memset(m->error_row_[c + kMaxComponents], 0.0,
cinfo->output_width * sizeof(float));
}
}
} // namespace jpegli