Merge pull request #101 from Moodstocks/png_generic

png.c: move apart non generic functions

4 files changed, 77 insertions, 49 deletions

diff --git a/generic/png.c b/generic/png.c
index 78f4426..24852de 100755
--- a/generic/png.c
+++ b/generic/png.c
@@ -10,46 +10,6 @@
  *
  * Clement: modified for Torch7.
  */
-#include <assert.h>
-
-/*
- * Bookkeeping struct for reading png data from memory
- */
-typedef struct {
-  unsigned char* buffer;
-  png_size_t offset;
-  png_size_t length;
-} libpng_(inmem_buffer);
-
-/*
- * Error message wrapper (single member struct to preserve `str` size info)
- */
-typedef struct {
-  char str[256];
-} libpng_(errmsg);
-
-/*
- * Call back for reading png data from memory
- */
-void libpng_(userReadData)(png_structp pngPtrSrc, png_bytep dest, png_size_t length) 
-{
-  libpng_(inmem_buffer)* src = png_get_io_ptr(pngPtrSrc);
-  assert(src->offset+length <= src->length);
-  memcpy(dest, src->buffer + src->offset, length);
-  src->offset += length;
-}
-
-/*
- * Custom error handling function (see `png_set_error_fn`)
- */
-static void libpng_(error_fn)(png_structp png_ptr, png_const_charp error_msg)
-{
-  libpng_(errmsg) *errmsg = png_get_error_ptr(png_ptr);
-  int max = sizeof(errmsg->str) - 1;
-  strncpy(errmsg->str, error_msg, max);
-  errmsg->str[max] = '\0';
-  longjmp(png_jmpbuf(png_ptr), 1);
-}
 
 static int libpng_(Main_load)(lua_State *L)
 {
@@ -64,8 +24,8 @@ static int libpng_(Main_load)(lua_State *L)
   png_bytep * row_pointers;
   size_t fread_ret;
   FILE* fp;
-  libpng_(inmem_buffer) inmem = {0};    /* source memory (if loading from memory) */
-  libpng_(errmsg) errmsg;
+  libpng_inmem_buffer inmem = {0};    /* source memory (if loading from memory) */
+  libpng_errmsg errmsg;
 
   const int load_from_file = luaL_checkint(L, 1);
 
@@ -96,7 +56,7 @@ static int libpng_(Main_load)(lua_State *L)
   if (!png_ptr)
     luaL_error(L, "[read_png] png_create_read_struct failed");
 
-  png_set_error_fn(png_ptr, &errmsg, libpng_(error_fn), NULL);
+  png_set_error_fn(png_ptr, &errmsg, libpng_error_fn, NULL);
 
   info_ptr = png_create_info_struct(png_ptr);
   if (!info_ptr)
@@ -109,7 +69,7 @@ static int libpng_(Main_load)(lua_State *L)
     png_init_io(png_ptr, fp);
   } else {
     /* set the read callback */
-    png_set_read_fn(png_ptr,(png_voidp)&inmem, libpng_(userReadData));
+    png_set_read_fn(png_ptr,(png_voidp)&inmem, libpng_userReadData);
   }
   png_set_sig_bytes(png_ptr, 8);
   png_read_info(png_ptr, info_ptr);
@@ -237,7 +197,7 @@ static int libpng_(Main_save)(lua_State *L)
   png_structp png_ptr;
   png_infop info_ptr;
   png_bytep * row_pointers;
-  libpng_(errmsg) errmsg;
+  libpng_errmsg errmsg;
 
   /* get dims and contiguous tensor */
   THTensor *tensorc = THTensor_(newContiguous)(tensor);
@@ -272,7 +232,7 @@ static int libpng_(Main_save)(lua_State *L)
   if (!png_ptr)
     luaL_error(L, "[write_png_file] png_create_write_struct failed");
 
-  png_set_error_fn(png_ptr, &errmsg, libpng_(error_fn), NULL);
+  png_set_error_fn(png_ptr, &errmsg, libpng_error_fn, NULL);
 
   info_ptr = png_create_info_struct(png_ptr);
   if (!info_ptr)
@@ -346,7 +306,7 @@ static int libpng_(Main_size)(lua_State *L)
 
   png_structp png_ptr;
   png_infop info_ptr;
-  libpng_(errmsg) errmsg;
+  libpng_errmsg errmsg;
   size_t fread_ret;
   /* open file and test for it being a png */
   FILE *fp = fopen(filename, "rb");
@@ -365,7 +325,7 @@ static int libpng_(Main_size)(lua_State *L)
   if (!png_ptr)
     luaL_error(L, "[get_png_size] png_create_read_struct failed");
 
-  png_set_error_fn(png_ptr, &errmsg, libpng_(error_fn), NULL);
+  png_set_error_fn(png_ptr, &errmsg, libpng_error_fn, NULL);
 
   info_ptr = png_create_info_struct(png_ptr);
   if (!info_ptr)
diff --git a/png.c b/png.c
index e139bfd..2034639 100644
--- a/png.c
+++ b/png.c
@@ -4,6 +4,7 @@
 #include <stdlib.h>
 #include <stdio.h>
 #include <string.h>
+#include <assert.h>
 
 #define PNG_DEBUG 3
 #include <png.h>
@@ -12,6 +13,47 @@
 #define torch_Tensor TH_CONCAT_STRING_3(torch., Real, Tensor)
 #define libpng_(NAME) TH_CONCAT_3(libpng_, Real, NAME)
 
+/*
+ * Bookkeeping struct for reading png data from memory
+ */
+typedef struct {
+  unsigned char* buffer;
+  png_size_t offset;
+  png_size_t length;
+} libpng_inmem_buffer;
+
+/*
+ * Call back for reading png data from memory
+ */
+static void
+libpng_userReadData(png_structp pngPtrSrc, png_bytep dest, png_size_t length)
+{
+  libpng_inmem_buffer* src = png_get_io_ptr(pngPtrSrc);
+  assert(src->offset+length <= src->length);
+  memcpy(dest, src->buffer + src->offset, length);
+  src->offset += length;
+}
+
+/*
+ * Error message wrapper (single member struct to preserve `str` size info)
+ */
+typedef struct {
+  char str[256];
+} libpng_errmsg;
+
+/*
+ * Custom error handling function (see `png_set_error_fn`)
+ */
+static void
+libpng_error_fn(png_structp png_ptr, png_const_charp error_msg)
+{
+  libpng_errmsg *errmsg = png_get_error_ptr(png_ptr);
+  int max = sizeof(errmsg->str) - 1;
+  strncpy(errmsg->str, error_msg, max);
+  errmsg->str[max] = '\0';
+  longjmp(png_jmpbuf(png_ptr), 1);
+}
+
 #include "generic/png.c"
 #include "THGenerateAllTypes.h"
 
diff --git a/test/corrupt-ihdr.png b/test/corrupt-ihdr.png
new file mode 100644
index 0000000..ca53ac9
--- /dev/null
+++ b/test/corrupt-ihdr.png
diff --git a/test/test_png.lua b/test/test_png.lua
index c01915d..376ef6f 100644
--- a/test/test_png.lua
+++ b/test/test_png.lua
@@ -5,7 +5,19 @@ local mytester = torch.Tester()
 local precision_mean = 1e-3
 local test = {}
 
-function checkPNG(imfile, depth, tensortype, want)
+local function toBlob(filename)
+  local f = torch.DiskFile(filename, 'r')
+  f:binary()
+  f:seekEnd()
+  local size = f:position() - 1
+  f:seek(1)
+  local blob = torch.ByteTensor(size)
+  f:readByte(blob:storage())
+  f:close()
+  return blob
+end
+
+local 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)
@@ -44,6 +56,20 @@ function test.LoadPNG()
   checkPNG('rgb16-2x1.png', 3, 'float', rgb16float)
 end
 
+function test.DecompressPNG()
+  mytester:assertTensorEq(
+    image.load('rgb2x1.png'),
+    image.decompressPNG(toBlob('rgb2x1.png')),
+    precision_mean,
+    'decompressed and loaded images should be equal'
+  )
+end
+
+function test.LoadCorruptedPNG()
+  local ok, _ = pcall(image.load, 'corrupt-ihdr.png')
+  mytester:assert(not ok, 'corrupted image should not be loaded')
+end
+
 -- Now run the test above
 mytester:add(test)
 mytester:run()