Kovid Goyal 8a5eebf409 ...
2016-05-09 09:30:36 +05:30

329 lines
13 KiB
C++

/*
* quantize.cpp
* Copyright (C) 2016 Kovid Goyal <kovid at kovidgoyal.net>
*
* octree based image quantization.
* See https://www.microsoft.com/msj/archive/S3F1.aspx for a simple to follow
* writeup on this algorithm
*
* Distributed under terms of the GPL3 license.
*/
#include <algorithm>
#include <cmath>
#include <stdio.h>
#include <QVector>
#include "imageops.h"
#ifdef _MSC_VER
typedef unsigned __int64 uint64_t;
typedef __int64 int64_t;
typedef unsigned __int32 uint32_t;
#define UINT64_MAX _UI64_MAX
#ifndef log2
static inline double log2(double x) { return log(x) / log((double)2) ; }
#endif
#else
#include <stdint.h>
#endif
#define MAX_DEPTH 8
#define MAX_COLORS 256
#define MAX(x, y) ((x) > (y)) ? (x) : (y)
#define MIN(x, y) ((x) < (y)) ? (x) : (y)
static const unsigned char BIT_MASK[8] = { 1 << 7, 1 << 6, 1 << 5, 1 << 4, 1 << 3, 1 << 2, 1 << 1, 1 };
static inline size_t get_index(const uint32_t r, const uint32_t g, const uint32_t b, const size_t level) {
return ((((r & BIT_MASK[level]) >> (7 - level)) << 2) | (((g & BIT_MASK[level]) >> (7 - level)) << 1) | ((b & BIT_MASK[level]) >> (7 - level)));
}
template <typename T> static inline T euclidean_distance(T r1, T g1, T b1, T r2, T g2, T b2) {
return r1 * r1 + r2 * r2 + g1 * g1 + g2 * g2 + b1 * b1 + b2 * b2 - 2 * (r1 * r2 + g1 * g2 + b1 * b2);
}
struct DoublePixel { double red; double green; double blue; };
template <class T> class Pool { // {{{
private:
QVector<T> nodes;
T *first_available;
public:
Pool<T>(size_t size) : nodes(size), first_available(nodes.data()) {
for (size_t i = 0; i < size - 1; i++) this->nodes[i].next_available_in_pool = &this->nodes[i+1];
}
T* checkout() {
T *ans = this->first_available;
if (ans == NULL) throw std::out_of_range("Something bad happened: ran out of nodes in the pool");
this->first_available = ans->next_available_in_pool;
if (this->first_available == NULL) {
// Grow the pool
int size = this->nodes.size();
this->nodes.resize(2*size);
this->first_available = &this->nodes[size];
for (int i = size; i < 2*size - 1; i++) this->nodes[i].next_available_in_pool = &this->nodes[i+1];
}
return ans;
}
void relinquish(T *node) {
node->reset();
node->next_available_in_pool = this->first_available;
this->first_available = node;
}
}; // }}}
class Node {
friend class Pool<Node>;
private:
bool is_leaf;
unsigned char index;
uint64_t pixel_count;
uint64_t red_sum;
uint64_t green_sum;
uint64_t blue_sum;
unsigned char red_avg;
unsigned char green_avg;
unsigned char blue_avg;
Node* next_reducible_node;
Node *next_available_in_pool;
Node* children[MAX_DEPTH];
public:
#ifdef _MSC_VER
// Disable the new behavior warning caused by children() below
#pragma warning( push )
#pragma warning (disable: 4351)
Node() : is_leaf(false), index(0), pixel_count(0), red_sum(0), green_sum(0), blue_sum(0), red_avg(0), green_avg(0), blue_avg(0), next_reducible_node(NULL), next_available_in_pool(NULL), children() {}
#pragma warning ( pop )
#endif
void reset() {
this->is_leaf = false;
this->pixel_count = 0;
this->red_sum = 0;
this->green_sum = 0;
this->blue_sum = 0;
this->next_reducible_node = NULL;
for (size_t i = 0; i < MAX_DEPTH; i++) this->children[i] = NULL;
}
void check_compiler() {
if (this->children[0] != NULL) throw std::runtime_error("Compiler failed to default initialize children");
}
inline Node* create_child(const size_t level, const size_t depth, unsigned int *leaf_count, Node **reducible_nodes, Pool<Node> &node_pool) {
Node *c = node_pool.checkout();
if (level == depth) {
c->is_leaf = true;
(*leaf_count)++;
} else {
c->next_reducible_node = reducible_nodes[level];
reducible_nodes[level] = c;
}
return c;
}
void add_color(const uint32_t r, const uint32_t g, const uint32_t b, const size_t depth, const size_t level, unsigned int *leaf_count, Node **reducible_nodes, Pool<Node> &node_pool) {
if (this->is_leaf) {
this->pixel_count++;
this->red_sum += r;
this->green_sum += g;
this->blue_sum += b;
} else {
size_t index = get_index(r, g, b, level);
if (this->children[index] == NULL) this->children[index] = this->create_child(level, depth, leaf_count, reducible_nodes, node_pool);
this->children[index]->add_color(r, g, b, depth, level + 1, leaf_count, reducible_nodes, node_pool);
}
}
void reduce(const size_t depth, unsigned int *leaf_count, Node **reducible_nodes, Pool<Node> &node_pool) {
size_t i = 0;
Node *node = NULL;
// Find the deepest level containing at least one reducible node
for (i=depth - 1; i > 0 && reducible_nodes[i] == NULL; i--);
// Reduce the node most recently added to the list at level i
// Could make this smarter by walking the linked list and choosing a
// node that has the least number of pixels or by storing error info
// on the nodes and using that
node = reducible_nodes[i];
reducible_nodes[i] = node->next_reducible_node;
for (i = 0; i < MAX_DEPTH; i++) {
if (node->children[i] != NULL) {
node->red_sum += node->children[i]->red_sum;
node->green_sum += node->children[i]->green_sum;
node->blue_sum += node->children[i]->blue_sum;
node->pixel_count += node->children[i]->pixel_count;
node_pool.relinquish(node->children[i]); node->children[i] = NULL;
(*leaf_count)--;
}
}
node->is_leaf = true; *leaf_count += 1;
}
void set_palette_colors(QRgb *color_table, unsigned char *index, bool compute_parent_averages) {
int i;
Node *child;
if (this->is_leaf) {
#define AVG_COLOR(x) ((unsigned char) ((double)this->x / (double)this->pixel_count))
this->red_avg = AVG_COLOR(red_sum); this->green_avg = AVG_COLOR(green_sum); this->blue_avg = AVG_COLOR(blue_sum);
color_table[*index] = qRgb(this->red_avg, this->green_avg, this->blue_avg);
this->index = (*index)++;
} else {
for (i = 0; i < MAX_DEPTH; i++) {
child = this->children[i];
if (child != NULL) {
child->set_palette_colors(color_table, index, compute_parent_averages);
if (compute_parent_averages) {
this->pixel_count += child->pixel_count;
this->red_sum += child->pixel_count * child->red_avg;
this->green_sum += child->pixel_count * child->green_avg;
this->blue_sum += child->pixel_count * child->blue_avg;
}
}
}
if (compute_parent_averages) {
this->red_avg = AVG_COLOR(red_sum); this->green_avg = AVG_COLOR(green_sum); this->blue_avg = AVG_COLOR(blue_sum);
}
}
}
unsigned char index_for_color(const uint32_t r, const uint32_t g, const uint32_t b, const size_t level) const {
if (this->is_leaf) return this->index;
size_t index = get_index(r, g, b, level);
if (this->children[index] == NULL) throw std::out_of_range("Something bad happened: could not follow tree for color");
return this->children[index]->index_for_color(r, g, b, level + 1);
}
unsigned char index_for_nearest_color(const uint32_t r, const uint32_t g, const uint32_t b, const size_t level) {
if (this->is_leaf) return this->index;
size_t index = get_index(r, g, b, level);
if (this->children[index] == NULL) {
uint64_t min_distance = UINT64_MAX, distance;
for(size_t i = 0; i < MAX_DEPTH; i++) {
Node *child = this->children[i];
if (child != NULL) {
distance = euclidean_distance<uint64_t>(r, g, b, child->red_avg, child->green_avg, child->blue_avg);
if (distance < min_distance) { min_distance = distance; index = i; }
}
}
}
return this->children[index]->index_for_nearest_color(r, g, b, level + 1);
}
};
// Image Dithering {{{
static inline void propagate_error(QVector<DoublePixel> &error_line, int c, unsigned char mult, DoublePixel &error) {
error_line[c].red += error.red * mult;
error_line[c].green += error.green * mult;
error_line[c].blue += error.blue * mult;
}
static inline QRgb apply_error(QRgb pixel, DoublePixel &error) {
#define AERR(w, i) MAX(0, MIN((int)(w(pixel) + error.i), 255))
return qRgb(AERR(qRed, red), AERR(qGreen, green), AERR(qBlue, blue));
}
static inline void calculate_error(QRgb new_pixel, QRgb old_pixel, DoublePixel &error) {
#define CERR(w) ((double)(w(old_pixel) - w(new_pixel)))/16.0
error.red = CERR(qRed);
error.green = CERR(qGreen);
error.blue = CERR(qBlue);
}
static void dither_image(const QImage &img, QImage &ans, QVector<QRgb> &color_table, Node &root) {
const QRgb *line = NULL;
QRgb pixel = 0, new_pixel = 0;
unsigned char *bits = NULL, index = 0;
int iheight = img.height(), iwidth = img.width(), r = 0, c = 0;
bool is_odd = false;
int start = 0, delta = 0;
DoublePixel error = {0, 0, 0};
const DoublePixel zero = {0, 0, 0};
QVector<DoublePixel> err1(iwidth), err2(iwidth), *line1 = NULL, *line2 = NULL;
for (r = 0; r < iheight; r++) {
line = reinterpret_cast<const QRgb*>(img.constScanLine(r));
bits = ans.scanLine(r);
is_odd = r & 1;
if (is_odd) { start = iwidth - 1; delta = -1; line1 = &err2; line2 = &err1; }
else { start = 0; delta = 1; line1 = &err1; line2 = &err2; }
line2->fill(zero);
for (c = start; 0 < (is_odd ? c + 1 : iwidth - c); c += delta) {
pixel = *(line + c);
new_pixel = apply_error(pixel, (*line1)[c]);
index = root.index_for_nearest_color(qRed(new_pixel), qGreen(new_pixel), qBlue(new_pixel), 0);
*(bits + c) = index;
calculate_error(color_table[index], pixel, error);
if (0 < (is_odd ? c : iwidth - c - 1)) {
propagate_error(*line1, c + delta, 7, error);
propagate_error(*line2, c + delta, 1, error);
}
propagate_error(*line2, c, 5, error);
if (0 < (is_odd ? iwidth - c - 1 : c)) propagate_error(*line2, c - delta, 3, error);
}
}
}
// }}}
QImage quantize(const QImage &image, unsigned int maximum_colors, bool dither) {
ScopedGILRelease PyGILRelease;
size_t depth = 0;
int iwidth = image.width(), iheight = image.height(), r, c;
QImage img(image), ans(iwidth, iheight, QImage::Format_Indexed8);
unsigned int leaf_count = 0;
unsigned char index = 0, *bits;
Node* reducible_nodes[MAX_DEPTH + 1] = {0};
Node root = Node();
QVector<QRgb> color_table = QVector<QRgb>(MAX_COLORS);
const QRgb* line = NULL;
root.check_compiler();
maximum_colors = MAX(2, MIN(MAX_COLORS, maximum_colors));
if (img.colorCount() > 0 && (size_t)img.colorCount() <= maximum_colors) return img; // Image is already quantized
if (img.hasAlphaChannel()) throw std::out_of_range("Cannot quantize image with transparency");
// TODO: Handle indexed image with colors > maximum_colors more efficiently
// by iterating over the color table rather than the pixels
if (img.format() != QImage::Format_RGB32) img = img.convertToFormat(QImage::Format_RGB32);
if (img.isNull()) throw std::bad_alloc();
// There can be at-most 8*(maximum_colors + 1) nodes, since we reduce the
// tree after each color is added Use an extra eight node just in case
// there is an off-by-one error somewhere :)
Pool<Node> node_pool((2 + maximum_colors) * 8);
depth = (size_t)log2(maximum_colors);
depth = MAX(2, MIN(depth, MAX_DEPTH));
for (r = 0; r < iheight; r++) {
line = reinterpret_cast<const QRgb*>(img.constScanLine(r));
for (c = 0; c < iwidth; c++) {
const QRgb pixel = *(line + c);
root.add_color(qRed(pixel), qGreen(pixel), qBlue(pixel), depth, 0, &leaf_count, reducible_nodes, node_pool);
while (leaf_count > maximum_colors)
root.reduce(depth, &leaf_count, reducible_nodes, node_pool);
}
}
if (leaf_count > maximum_colors) throw std::out_of_range("Leaf count > max colors, something bad happened");
color_table.resize(leaf_count);
root.set_palette_colors(color_table.data(), &index, dither);
ans.setColorTable(color_table);
if (dither) {
dither_image(img, ans, color_table, root);
} else {
for (r = 0; r < iheight; r++) {
line = reinterpret_cast<const QRgb*>(img.constScanLine(r));
bits = ans.scanLine(r);
for (c = 0; c < iwidth; c++) {
const QRgb pixel = *(line + c);
*(bits + c) = root.index_for_color(qRed(pixel), qGreen(pixel), qBlue(pixel), 0);
}
}
}
return ans;
}