From c4c0be53e8151a61529822c1c11a99fad3f6ec20 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Jos=C3=A9=20Fonseca?= <jose.r.fonseca@gmail.com>
Date: Tue, 24 May 2011 09:49:23 +0100
Subject: [PATCH] Handle an arbitrary number of image channels.

---
 glstate.cpp | 11 +++++------
 image.cpp   | 36 +++++++++++++++++++++++++++++++-----
 image.hpp   | 12 +++++++-----
 3 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/glstate.cpp b/glstate.cpp
index 2fa4cb5..71aad7d 100644
--- a/glstate.cpp
+++ b/glstate.cpp
@@ -407,18 +407,17 @@ getDrawableBounds(GLint *width, GLint *height) {
 
 Image::Image *
 getDrawBufferImage(GLenum format) {
-    GLint width, height;
-
-    if (format != GL_RGBA) {
-        // FIXME: this function can only handle 4-channel images
+    GLint channels = __gl_format_channels(format);
+    if (channels > 4) {
         return NULL;
     }
 
+    GLint width, height;
     if (!getDrawableBounds(&width, &height)) {
         return NULL;
     }
 
-    Image::Image *image = new Image::Image(width, height, true);
+    Image::Image *image = new Image::Image(width, height, channels, true);
     if (!image) {
         return NULL;
     }
@@ -439,7 +438,7 @@ getDrawBufferImage(GLenum format) {
     glPixelStorei(GL_PACK_SKIP_IMAGES, 0);
     glPixelStorei(GL_PACK_ALIGNMENT, 1);
 
-    glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE, image->pixels);
+    glReadPixels(0, 0, width, height, format, GL_UNSIGNED_BYTE, image->pixels);
 
     glPopClientAttrib();
 
diff --git a/image.cpp b/image.cpp
index 922c4d4..3aa0fe5 100644
--- a/image.cpp
+++ b/image.cpp
@@ -1,5 +1,6 @@
 /**************************************************************************
  *
+ * Copyright 2011 Jose Fonseca
  * Copyright 2008-2010 VMware, Inc.
  * All Rights Reserved.
  *
@@ -27,6 +28,7 @@
 #include <zlib.h>
 #include <png.h>
 
+#include <assert.h>
 #include <math.h>
 #include <stdint.h>
 #include <stdlib.h>
@@ -73,6 +75,8 @@ struct Pixel {
 
 bool
 Image::writeBMP(const char *filename) const {
+    assert(channels == 4);
+
     struct FileHeader bmfh;
     struct InfoHeader bmih;
     unsigned x, y;
@@ -164,7 +168,26 @@ Image::writePNG(const char *filename) const {
 
     png_init_io(png_ptr, fp);
 
-    png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_RGB_ALPHA,
+    int color_type;
+    switch (channels) {
+    case 4:
+        color_type = PNG_COLOR_TYPE_RGB_ALPHA;
+        break;
+    case 3:
+        color_type = PNG_COLOR_TYPE_RGB;
+        break;
+    case 2:
+        color_type = PNG_COLOR_TYPE_GRAY_ALPHA;
+        break;
+    case 1:
+        color_type = PNG_COLOR_TYPE_GRAY;
+        break;
+    default:
+        assert(0);
+        return false;
+    }
+
+    png_set_IHDR(png_ptr, info_ptr, width, height, 8, color_type,
         PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
 
     png_set_compression_level(png_ptr, Z_DEFAULT_COMPRESSION);
@@ -173,13 +196,13 @@ Image::writePNG(const char *filename) const {
 
     if (!flipped) {
         for (unsigned y = 0; y < height; ++y) {
-            png_bytep row = (png_bytep)(pixels + y*width*4);
+            png_bytep row = (png_bytep)(pixels + y*width*channels);
             png_write_rows(png_ptr, &row, 1);
         }
     } else {
         unsigned y = height;
         while (y--) {
-            png_bytep row = (png_bytep)(pixels + y*width*4);
+            png_bytep row = (png_bytep)(pixels + y*width*channels);
             png_write_rows(png_ptr, &row, 1);
         }
     }
@@ -279,20 +302,23 @@ no_fp:
 double Image::compare(Image &ref)
 {
     if (width != ref.width ||
-        height != ref.height) {
+        height != ref.height ||
+        channels != ref.channels) {
         return 0.0;
     }
 
     const unsigned char *pSrc = start();
     const unsigned char *pRef = ref.start();
 
+    assert(channels >= 3);
+
     unsigned long long error = 0;
     for (unsigned y = 0; y < height; ++y) {
         for (unsigned  x = 0; x < width; ++x) {
             // FIXME: Ignore alpha channel until we are able to pick a visual
             // that matches the traces
             for (unsigned  c = 0; c < 3; ++c) {
-                int delta = pSrc[x*4 + c] - pRef[x*4 + c];
+                int delta = pSrc[x*channels + c] - pRef[x*channels + c];
                 error += delta*delta;
             }
         }
diff --git a/image.hpp b/image.hpp
index f46fb31..e02abed 100644
--- a/image.hpp
+++ b/image.hpp
@@ -41,6 +41,7 @@ class Image {
 public:
     unsigned width;
     unsigned height;
+    unsigned channels;
 
     // Flipped vertically or not
     bool flipped;
@@ -48,11 +49,12 @@ public:
     // Pixels in RGBA format
     unsigned char *pixels;
 
-    inline Image(unsigned w, unsigned h, bool f = false) : 
+    inline Image(unsigned w, unsigned h, unsigned c = 4, bool f = false) : 
         width(w),
         height(h),
+        channels(c),
         flipped(f),
-        pixels(new unsigned char[h*w*4])
+        pixels(new unsigned char[h*w*c])
     {}
 
     inline ~Image() {
@@ -60,15 +62,15 @@ public:
     }
 
     inline unsigned char *start(void) {
-        return flipped ? pixels + (height - 1)*width*4 : pixels;
+        return flipped ? pixels + (height - 1)*width*channels : pixels;
     }
 
     inline unsigned char *end(void) {
-        return flipped ? pixels - width*4 : pixels + height*width*4;
+        return flipped ? pixels - width*channels : pixels + height*width*channels;
     }
 
     inline signed stride(void) const {
-        return flipped ? -width*4 : width*4;
+        return flipped ? -width*channels : width*channels;
     }
 
     bool writeBMP(const char *filename) const;
-- 
2.45.2