Merge pull request #76 from krasin/two-args-png-load

png: add support for 16-bit PNG files.

7 files changed, 89 insertions, 14 deletions

diff --git a/generic/png.c b/generic/png.c
index 1e34355..6702c12 100755
--- a/generic/png.c
+++ b/generic/png.c
@@ -37,7 +37,7 @@ static int libpng_(Main_load)(lua_State *L)
 
   png_byte header[8];    // 8 is the maximum size that can be checked
 
-  int width, height;
+  int width, height, bit_depth;
   png_byte color_type;
   
   png_structp png_ptr;
@@ -95,6 +95,7 @@ static int libpng_(Main_load)(lua_State *L)
   width      = png_get_image_width(png_ptr, info_ptr);
   height     = png_get_image_height(png_ptr, info_ptr);
   color_type = png_get_color_type(png_ptr, info_ptr);
+  bit_depth  = png_get_bit_depth(png_ptr, info_ptr);
   png_read_update_info(png_ptr, info_ptr);
 
   /* get depth */
@@ -105,7 +106,7 @@ static int libpng_(Main_load)(lua_State *L)
     depth = 3;
   else if (color_type == PNG_COLOR_TYPE_GRAY)
   {
-    if(png_get_bit_depth(png_ptr, info_ptr) < 8)
+    if(bit_depth < 8)
     {
       png_set_expand_gray_1_2_4_to_8(png_ptr);
       png_read_update_info(png_ptr, info_ptr);
@@ -123,7 +124,7 @@ static int libpng_(Main_load)(lua_State *L)
   else
     luaL_error(L, "[read_png_file] Unknown color space");
 
-  if(png_get_bit_depth(png_ptr, info_ptr) < 8)
+  if(bit_depth < 8)
   {
     png_set_strip_16(png_ptr);
     png_read_update_info(png_ptr, info_ptr);
@@ -148,13 +149,31 @@ static int libpng_(Main_load)(lua_State *L)
 
   /* convert image to dest tensor */
   int x,k;
-  for (k=0; k<depth; k++) {
-    for (y=0; y<height; y++) {
-      png_byte* row = row_pointers[y];
-      for (x=0; x<width; x++) {
-        *tensor_data++ = (real)row[x*depth+k];
-        //png_byte val = row[x*depth+k];
-        //THTensor_(set3d)(tensor, k, y, x, (real)val);
+  if ((bit_depth == 16) && (sizeof(real) > 1)) {
+    for (k=0; k<depth; k++) {
+      for (y=0; y<height; y++) {
+	png_byte* row = row_pointers[y];
+	for (x=0; x<width; x++) {
+	  // PNG is big-endian
+	  int val = ((int)row[(x*depth+k)*2] << 8) + row[(x*depth+k)*2+1];
+	  *tensor_data++ = (real)val;
+	}
+      }
+    }
+  } else {
+    int stride = 1;
+    if (bit_depth == 16) {
+      /* PNG has 16 bit color depth, but the tensor type is byte. */
+      stride = 2;
+    }
+    for (k=0; k<depth; k++) {
+      for (y=0; y<height; y++) {
+	png_byte* row = row_pointers[y];
+	for (x=0; x<width; x++) {
+	  *tensor_data++ = (real)row[(x*depth+k)*stride];
+	  //png_byte val = row[x*depth+k];
+	  //THTensor_(set3d)(tensor, k, y, x, (real)val);
+	}
       }
     }
   }
@@ -176,7 +195,13 @@ static int libpng_(Main_load)(lua_State *L)
 
   /* return tensor */
   luaT_pushudata(L, tensor, torch_Tensor);
-  return 1;
+
+  if (bit_depth < 8) {
+    bit_depth = 8;
+  }
+  lua_pushnumber(L, bit_depth);
+
+  return 2;
 }
 
 static int libpng_(Main_save)(lua_State *L)
diff --git a/init.lua b/init.lua
index b19abf9..0819ec0 100755
--- a/init.lua
+++ b/init.lua
@@ -118,8 +118,9 @@ local function decompress(tensor, depth, tensortype)
 end
 rawset(image, 'decompress', decompress)
 
-local function processPNG(img, depth, tensortype)
+local function processPNG(img, depth, bit_depth, tensortype)
     local MAXVAL = 255
+    if bit_depth == 16 then MAXVAL = 65535 end
     if tensortype ~= 'byte' then
         img:mul(1/MAXVAL)
     end
@@ -132,8 +133,8 @@ local function loadPNG(filename, depth, tensortype)
       dok.error('libpng package not found, please install libpng','image.loadPNG')
    end
    local load_from_file = 1
-   local a = template(tensortype).libpng.load(load_from_file, filename)
-   return processPNG(a, depth, tensortype)
+   local a, bit_depth = template(tensortype).libpng.load(load_from_file, filename)
+   return processPNG(a, depth, bit_depth, tensortype)
 end
 rawset(image, 'loadPNG', loadPNG)
 
diff --git a/test/gray16-1x2.png b/test/gray16-1x2.png
new file mode 100644
index 0000000..9b3cb5e
--- /dev/null
+++ b/test/gray16-1x2.png
diff --git a/test/gray3x1.png b/test/gray3x1.png
new file mode 100644
index 0000000..ce89719
--- /dev/null
+++ b/test/gray3x1.png
diff --git a/test/rgb16-2x1.png b/test/rgb16-2x1.png
new file mode 100644
index 0000000..3aab682
--- /dev/null
+++ b/test/rgb16-2x1.png
diff --git a/test/rgb2x1.png b/test/rgb2x1.png
new file mode 100644
index 0000000..60a0e9f
--- /dev/null
+++ b/test/rgb2x1.png
diff --git a/test/test_png.lua b/test/test_png.lua
new file mode 100644
index 0000000..c01915d
--- /dev/null
+++ b/test/test_png.lua
@@ -0,0 +1,49 @@
+require 'image'
+
+-- Create an instance of the test framework
+local mytester = torch.Tester()
+local precision_mean = 1e-3
+local test = {}
+
+function checkPNG(imfile, depth, tensortype, want)
+  local img = image.load(imfile, depth, tensortype)
+  -- Tensors have to be converted to double, since assertTensorEq does not support ByteTensor
+  print('img: ', img)
+  print('want: ', want)
+  mytester:assertTensorEq(img:double(), want:double(), precision_mean,
+                          string.format('%s: pixel values are unexpected', imfile))
+end
+
+function test.LoadPNG()
+  -- Gray 8-bit PNG image with width = 3, height = 1
+  local gray8byte = torch.ByteTensor({{{0,127,255}}})
+  checkPNG('gray3x1.png', 1, 'byte', gray8byte)
+
+  local gray8double = torch.DoubleTensor({{{0, 127/255, 1}}})
+  checkPNG('gray3x1.png', 1, 'double', gray8double)
+
+  -- Gray 16-bit PNG image with width=1, height = 2
+  local gray16byte = torch.ByteTensor({{{0, 255}}})
+  checkPNG('gray16-1x2.png', 1, 'byte', gray16byte)
+
+  local gray16float = torch.FloatTensor({{{0, 65534/65535}}})
+  checkPNG('gray16-1x2.png', 1, 'float', gray16float)
+
+  -- Color 8-bit PNG image with width = 2, height = 1
+  local rgb8byte = torch.ByteTensor({{{255, 0, 0, 127, 63, 0}}})
+  checkPNG('rgb2x1.png', 3, 'byte', rgb8byte)
+
+  local rgb8float = torch.FloatTensor({{{1, 0, 0, 127/255, 63/255, 0}}})
+  checkPNG('rgb2x1.png', 3, 'float', rgb8float)
+
+  -- Color 16-bit PNG image with width = 2, height = 1
+  local rgb16byte = torch.ByteTensor({{{255, 0, 0, 127, 63, 0}}})
+  checkPNG('rgb16-2x1.png', 3, 'byte', rgb16byte)
+
+  local rgb16float = torch.FloatTensor({{{1, 0, 0, 32767/65535, 16383/65535, 0}}})
+  checkPNG('rgb16-2x1.png', 3, 'float', rgb16float)
+end
+
+-- Now run the test above
+mytester:add(test)
+mytester:run()