pnmimage: add quantize feature to reduce number of colors in image

panda/src/pnmimage/pnmImage.I

@@ -71,6 +71,57 @@ clamp_val(int input_value) const {
   return (xelval)std::min(std::max(0, input_value), (int)get_maxval());
 }
 
+/**
+ * A handy function to scale non-alpha values from [0..1] to
+ * [0..get_maxval()].  Do not use this for alpha values, see to_alpha_val.
+ */
+INLINE xel PNMImage::
+to_val(const LRGBColorf &value) const {
+  xel col;
+  switch (_xel_encoding) {
+  case XE_generic:
+  case XE_generic_alpha:
+    {
+      LRGBColorf scaled = value * get_maxval() + 0.5f;
+      col.r = clamp_val((int)scaled[0]);
+      col.g = clamp_val((int)scaled[1]);
+      col.b = clamp_val((int)scaled[2]);
+    }
+    break;
+
+  case XE_generic_sRGB:
+  case XE_generic_sRGB_alpha:
+    col.r = clamp_val((int)
+      (encode_sRGB_float(value[0]) * get_maxval() + 0.5f));
+    col.g = clamp_val((int)
+      (encode_sRGB_float(value[1]) * get_maxval() + 0.5f));
+    col.b = clamp_val((int)
+      (encode_sRGB_float(value[2]) * get_maxval() + 0.5f));
+    break;
+
+  case XE_uchar_sRGB:
+  case XE_uchar_sRGB_alpha:
+    encode_sRGB_uchar(LColorf(value, 0.0f), col);
+    break;
+
+  case XE_uchar_sRGB_sse2:
+  case XE_uchar_sRGB_alpha_sse2:
+    encode_sRGB_uchar_sse2(LColorf(value, 0.0f), col);
+    break;
+
+  case XE_scRGB:
+  case XE_scRGB_alpha:
+    {
+      LRGBColorf scaled = value * 8192.f + 4096.5f;
+      col.r = std::min(std::max(0, (int)scaled[0]), 65535);
+      col.g = std::min(std::max(0, (int)scaled[1]), 65535);
+      col.b = std::min(std::max(0, (int)scaled[2]), 65535);
+    }
+    break;
+  }
+  return col;
+}
+
 /**
  * A handy function to scale non-alpha values from [0..1] to
  * [0..get_maxval()].  Do not use this for alpha values, see to_alpha_val.
@@ -112,6 +163,44 @@ to_alpha_val(float input_value) const {
   return clamp_val((int)(input_value * get_maxval() + 0.5));
 }
 
+/**
+ * A handy function to scale non-alpha values from [0..get_maxval()] to
+ * [0..1].  Do not use this for alpha values, see from_alpha_val.
+ */
+INLINE LRGBColorf PNMImage::
+from_val(const xel &col) const {
+  switch (_xel_encoding) {
+  case XE_generic:
+  case XE_generic_alpha:
+    return LRGBColorf(col.r, col.g, col.b) * _inv_maxval;
+
+  case XE_generic_sRGB:
+  case XE_generic_sRGB_alpha:
+    return LRGBColorf(
+      decode_sRGB_float(col.r * _inv_maxval),
+      decode_sRGB_float(col.g * _inv_maxval),
+      decode_sRGB_float(col.b * _inv_maxval));
+
+  case XE_uchar_sRGB:
+  case XE_uchar_sRGB_alpha:
+  case XE_uchar_sRGB_sse2:
+  case XE_uchar_sRGB_alpha_sse2:
+    return LRGBColorf(
+      decode_sRGB_float((unsigned char)col.r),
+      decode_sRGB_float((unsigned char)col.g),
+      decode_sRGB_float((unsigned char)col.b));
+
+  case XE_scRGB:
+  case XE_scRGB_alpha:
+    return LRGBColorf((int)col.r - 4096,
+                      (int)col.g - 4096,
+                      (int)col.b - 4096) * (1.f / 8192.f);
+
+  default:
+    return LRGBColorf(0);
+  }
+}
+
 /**
  * A handy function to scale non-alpha values from [0..get_maxval()] to
  * [0..1].  Do not use this for alpha values, see from_alpha_val.
@@ -479,39 +568,7 @@ set_alpha_val(int x, int y, xelval a) {
 INLINE LRGBColorf PNMImage::
 get_xel(int x, int y) const {
   nassertr(x >= 0 && x < _x_size && y >= 0 && y < _y_size, LRGBColorf::zero());
-
+  return from_val(row(y)[x]);
-  const xel &col = row(y)[x];
-
-  switch (_xel_encoding) {
-  case XE_generic:
-  case XE_generic_alpha:
-    return LRGBColorf(col.r, col.g, col.b) * _inv_maxval;
-
-  case XE_generic_sRGB:
-  case XE_generic_sRGB_alpha:
-    return LRGBColorf(
-      decode_sRGB_float(col.r * _inv_maxval),
-      decode_sRGB_float(col.g * _inv_maxval),
-      decode_sRGB_float(col.b * _inv_maxval));
-
-  case XE_uchar_sRGB:
-  case XE_uchar_sRGB_alpha:
-  case XE_uchar_sRGB_sse2:
-  case XE_uchar_sRGB_alpha_sse2:
-    return LRGBColorf(
-      decode_sRGB_float((unsigned char)col.r),
-      decode_sRGB_float((unsigned char)col.g),
-      decode_sRGB_float((unsigned char)col.b));
-
-  case XE_scRGB:
-  case XE_scRGB_alpha:
-    return LRGBColorf((int)col.r - 4096,
-                      (int)col.g - 4096,
-                      (int)col.b - 4096) * (1.f / 8192.f);
-
-  default:
-    return LRGBColorf(0);
-  }
 }
 
 /**
@@ -521,50 +578,7 @@ get_xel(int x, int y) const {
 INLINE void PNMImage::
 set_xel(int x, int y, const LRGBColorf &value) {
   nassertv(x >= 0 && x < _x_size && y >= 0 && y < _y_size);
-
+  row(y)[x] = to_val(value);
-  xel &col = row(y)[x];
-
-  switch (_xel_encoding) {
-  case XE_generic:
-  case XE_generic_alpha:
-    {
-      LRGBColorf scaled = value * get_maxval() + 0.5f;
-      col.r = clamp_val((int)scaled[0]);
-      col.g = clamp_val((int)scaled[1]);
-      col.b = clamp_val((int)scaled[2]);
-    }
-    break;
-
-  case XE_generic_sRGB:
-  case XE_generic_sRGB_alpha:
-    col.r = clamp_val((int)
-      (encode_sRGB_float(value[0]) * get_maxval() + 0.5f));
-    col.g = clamp_val((int)
-      (encode_sRGB_float(value[1]) * get_maxval() + 0.5f));
-    col.b = clamp_val((int)
-      (encode_sRGB_float(value[2]) * get_maxval() + 0.5f));
-    break;
-
-  case XE_uchar_sRGB:
-  case XE_uchar_sRGB_alpha:
-    encode_sRGB_uchar(LColorf(value, 0.0f), col);
-    break;
-
-  case XE_uchar_sRGB_sse2:
-  case XE_uchar_sRGB_alpha_sse2:
-    encode_sRGB_uchar_sse2(LColorf(value, 0.0f), col);
-    break;
-
-  case XE_scRGB:
-  case XE_scRGB_alpha:
-    {
-      LRGBColorf scaled = value * 8192.f + 4096.5f;
-      col.r = std::min(std::max(0, (int)scaled[0]), 65535);
-      col.g = std::min(std::max(0, (int)scaled[1]), 65535);
-      col.b = std::min(std::max(0, (int)scaled[2]), 65535);
-    }
-    break;
-  }
 }
 
 /**

panda/src/pnmimage/pnmImage.cxx

@@ -1928,6 +1928,51 @@ make_histogram(PNMImage::Histogram &histogram) {
   histogram.swap(pixels, hist_map);
 }
 
+/**
+ * Reduces the number of unique colors in the image to (at most) the given
+ * count.  Fewer colors than requested may be left in the image after this
+ * operation, but never more.
+ *
+ * At present, this is only supported on images without an alpha channel.
+ *
+ * @since 1.10.5
+ */
+void PNMImage::
+quantize(size_t max_colors) {
+  nassertv(_array != nullptr);
+  nassertv(!has_alpha());
+  size_t array_size = _x_size * _y_size;
+
+  // Get all the unique colors in this image.
+  pmap<xel, xel> color_map;
+  for (size_t i = 0; i < array_size; ++i) {
+    color_map[_array[i]];
+  }
+
+  size_t num_colors = color_map.size();
+  if (num_colors <= max_colors) {
+    // We are already down to the requested number of colors.
+    return;
+  }
+
+  // Collect all the colors into a contiguous array.
+  xel *colors = (xel *)alloca(num_colors * sizeof(xel));
+  size_t i = 0;
+  for (pmap<xel, xel>::const_iterator it = color_map.begin();
+       it != color_map.end(); ++it) {
+    colors[i++] = it->first;
+  }
+  nassertv(i == num_colors);
+
+  // Apply the median cut algorithm, which will give us a color map.
+  r_quantize(color_map, max_colors, colors, num_colors);
+
+  // Replace all the existing colors with the corresponding bucket average.
+  for (size_t i = 0; i < array_size; ++i) {
+    _array[i] = color_map[_array[i]];
+  }
+}
+
 /**
  * Fills the image with a grayscale perlin noise pattern based on the
  * indicated parameters.  Uses set_xel to set the grayscale values.  The sx
@@ -2161,6 +2206,92 @@ setup_encoding() {
   }
 }
 
+/**
+ * Recursive implementation of quantize() using the median cut algorithm.
+ */
+void PNMImage::
+r_quantize(pmap<xel, xel> &color_map, size_t max_colors,
+           xel *colors, size_t num_colors) {
+  if (num_colors <= max_colors) {
+    // All points in this bucket can be preserved 1:1.
+    for (size_t i = 0; i < num_colors; ++i) {
+      const xel &col = colors[i];
+      color_map[col] = col;
+    }
+    return;
+  }
+  else if (max_colors == 1) {
+    // We've reached the target.  Calculate the average, in linear space.
+    LRGBColorf avg(0);
+    for (size_t i = 0; i < num_colors; ++i) {
+      avg += from_val(colors[i]);
+    }
+    avg *= 1.0f / num_colors;
+    xel avg_val = to_val(avg);
+
+    // Map all colors in this bucket to the avg.
+    for (size_t i = 0; i < num_colors; ++i) {
+      color_map[colors[i]] = avg_val;
+    }
+    return;
+  }
+  else if (max_colors == 0) {
+    // Not sure how this happens, but we can't preserve any color here.
+    return;
+  }
+
+  // Find the minimum/maximum RGB values.  We should probably do this in
+  // linear space, but eh.
+  xelval min_r = _maxval;
+  xelval min_g = _maxval;
+  xelval min_b = _maxval;
+  xelval max_r = 0, max_g = 0, max_b = 0;
+  for (size_t i = 0; i < num_colors; ++i) {
+    const xel &col = colors[i];
+    min_r = std::min(min_r, col.r);
+    max_r = std::max(max_r, col.r);
+    min_g = std::min(min_g, col.g);
+    max_g = std::max(max_g, col.g);
+    min_b = std::min(min_b, col.b);
+    max_b = std::max(max_b, col.b);
+  }
+
+  int diff_r = max_r - min_r;
+  int diff_g = max_g - min_g;
+  int diff_b = max_b - min_b;
+
+  auto sort_by_red = [](const xel &c1, const xel &c2) {
+    return c1.r < c2.r;
+  };
+  auto sort_by_green = [](const xel &c1, const xel &c2) {
+    return c1.g < c2.g;
+  };
+  auto sort_by_blue = [](const xel &c1, const xel &c2) {
+    return c1.b < c2.b;
+  };
+
+  // Sort by the component with the most variation.
+  if (diff_g >= diff_r) {
+    if (diff_g >= diff_b) {
+      std::sort(colors, colors + num_colors, sort_by_green);
+    } else {
+      std::sort(colors, colors + num_colors, sort_by_blue);
+    }
+  } else if (diff_r >= diff_b) {
+    std::sort(colors, colors + num_colors, sort_by_red);
+  } else {
+    std::sort(colors, colors + num_colors, sort_by_blue);
+  }
+
+  // Subdivide the sorted colors into two buckets, and recurse.
+  size_t max_colors_1 = max_colors / 2;
+  size_t max_colors_2 = max_colors - max_colors_1;
+  size_t num_colors_1 = num_colors / 2;
+  size_t num_colors_2 = num_colors - num_colors_1;
+  r_quantize(color_map, max_colors_1, colors, num_colors_1);
+  r_quantize(color_map, max_colors_2, colors + num_colors_1, num_colors_2);
+}
+
 /**
  * Recursively fills in the minimum distance measured from a certain set of
  * points into the gray channel.

panda/src/pnmimage/pnmImage.h

@@ -68,8 +68,10 @@ PUBLISHED:
   INLINE ~PNMImage();
 
   INLINE xelval clamp_val(int input_value) const;
+  INLINE xel to_val(const LRGBColorf &input_value) const;
   INLINE xelval to_val(float input_value) const;
   INLINE xelval to_alpha_val(float input_value) const;
+  INLINE LRGBColorf from_val(const xel &input_value) const;
   INLINE float from_val(xelval input_value) const;
   INLINE float from_alpha_val(xelval input_value) const;
 
@@ -254,6 +256,7 @@ PUBLISHED:
                                   int xborder = 0, int yborder = 0);
 
   void make_histogram(Histogram &hist);
+  void quantize(size_t max_colors);
   BLOCKING void perlin_noise_fill(float sx, float sy, int table_size = 256,
                                   unsigned long seed = 0);
   void perlin_noise_fill(StackedPerlinNoise2 &perlin);
@@ -346,6 +349,9 @@ private:
   void setup_rc();
   void setup_encoding();
 
+  void r_quantize(pmap<xel, xel> &color_map, size_t max_colors,
+                  xel *colors, size_t num_colors);
+
 PUBLISHED:
   PNMImage operator ~() const;
 

panda/src/pnmimage/pnmimage_base.h

@@ -59,6 +59,22 @@ PUBLISHED:
   void operator *= (const double mult)
     { r *= mult; g *= mult; b *= mult; }
 
+  bool operator == (const pixel &other) {
+    return r == other.r && g == other.g && r == other.r;
+  }
+  bool operator != (const pixel &other) {
+    return r != other.r || g != other.g || r != other.r;
+  }
+  bool operator < (const pixel &other) const {
+    if (r != other.r) {
+      return r < other.r;
+    }
+    if (g != other.g) {
+      return g < other.g;
+    }
+    return b < other.b;
+  }
+
 #ifdef HAVE_PYTHON
   static int size() { return 3; }
   void output(std::ostream &out) {

tests/pnmimage/test_pnmimage.py

@@ -1,4 +1,5 @@
 from panda3d.core import PNMImage, PNMImageHeader
+from random import randint
 
 
 def test_pixelspec_ctor():
@@ -19,3 +20,53 @@ def test_pixelspec_coerce():
     img = PNMImage(1, 1, 4)
     img.set_pixel(0, 0, (1, 2, 3, 4))
     assert img.get_pixel(0, 0) == (1, 2, 3, 4)
+
+
+def test_pnmimage_to_val():
+    img = PNMImage(1, 1)
+    assert img.to_val(-0.5) == 0
+    assert img.to_val(0.0) == 0
+    assert img.to_val(0.5) == 128
+    assert img.to_val(1.0) == 255
+    assert img.to_val(2.0) == 255
+
+
+def test_pnmimage_from_val():
+    img = PNMImage(1, 1)
+    assert img.from_val(0) == 0.0
+    assert img.to_val(img.from_val(128)) == 128
+    assert img.from_val(255) == 1.0
+
+
+def test_pnmimage_quantize():
+    img = PNMImage(32, 32, 3)
+
+    for x in range(32):
+        for y in range(32):
+            img.set_xel_val(x, y, randint(0, 100), randint(50, 100), randint(0, 1))
+
+    hist = PNMImage.Histogram()
+    img.make_histogram(hist)
+    num_colors = hist.get_num_pixels()
+    assert num_colors > 100
+
+    img2 = PNMImage(img)
+    img2.quantize(100)
+    hist = PNMImage.Histogram()
+    img2.make_histogram(hist)
+    assert hist.get_num_pixels() <= 100
+
+    # Make sure that this is reasonably close
+    max_dist = 0
+    for x in range(32):
+        for y in range(32):
+            diff = img.get_xel(x, y) - img2.get_xel(x, y)
+            max_dist = max(max_dist, diff.length_squared())
+
+            # Also make sure that they are not out of range of the original
+            col = img2.get_xel_val(x, y)
+            assert col.r <= 100
+            assert col.g >= 50 and col.g <= 100
+            assert col.b in (0, 1)
+
+    assert max_dist < 0.1 ** 2