Merge pull request #52 from jonathantompson/compressJPG

Added docs for decompressJPG.  Added compressJPG function.

4 files changed, 138 insertions, 8 deletions

diff --git a/README.md b/README.md
index 741e67b..456c601 100644
--- a/README.md
+++ b/README.md
@@ -39,6 +39,31 @@ Saves Tensor `tensor` to disk at path `filename`. The format to which
 the image is saved is extrapolated from the `filename`'s extension suffix.
 The `tensor` should be of size `nChannel x height x width`.
 
+<a name="image.decompressJPG"/>
+### [res] image.decompressJPG(tensor, [depth, tensortype]) ###
+Decompresses an image from a ByteTensor in memory having `depth` channels (1 or 3)
+into a [Tensor](https://github.com/torch/torch7/blob/master/doc/tensor.md#tensor)
+of type `tensortype` (*float*, *double* or *byte*). The last two arguments
+are optional.
+
+Usage:
+```lua
+local fin = torch.DiskFile(imfile, 'r')
+fin:binary()
+fin:seekEnd()
+local file_size_bytes = fin:position() - 1
+fin:seek(1)
+local img_binary = torch.ByteTensor(file_size_bytes)
+fin:readByte(img_binary:storage())
+fin:close()
+-- Then when you're ready to decompress the ByteTensor:
+im = image.decompressJPG(img_binary)
+```
+
+<a name="image.compressJPG"/>
+### [res] image.compressJPG(tensor, [quality]) ###
+Compresses an image to a ByteTensor in memory.  Optional quality is between 1 and 100 and adjusts compression quality.
+
 <a name="image.simpletrans"/>
 ## Simple Transformations ##
 This section includes simple but very common image transformations 
diff --git a/generic/jpeg.c b/generic/jpeg.c
index 099486e..fa00fff 100755
--- a/generic/jpeg.c
+++ b/generic/jpeg.c
@@ -333,12 +333,27 @@ static int libjpeg_(Main_load)(lua_State *L)
  *
  */
 int libjpeg_(Main_save)(lua_State *L) {
+  unsigned char *inmem = NULL;  /* destination memory (if saving to memory) */
+  unsigned long inmem_size = 0;  /* destination memory size (bytes) */
+
   /* get args */
   const char *filename = luaL_checkstring(L, 1);
   THTensor *tensor = luaT_checkudata(L, 2, torch_Tensor);  
   THTensor *tensorc = THTensor_(newContiguous)(tensor);
   real *tensor_data = THTensor_(data)(tensorc);
 
+  const int save_to_file = luaL_checkint(L, 3);
+
+  THByteTensor* tensor_dest = NULL;
+  if (save_to_file == 0) {
+    tensor_dest = luaT_checkudata(L, 5, "torch.ByteTensor");
+  }
+
+  int quality = luaL_checkint(L, 4);
+  if (quality < 0 || quality > 100) {
+    luaL_error(L, "quality should be between 0 and 100");
+  }
+
   /* jpeg struct */
   struct jpeg_compress_struct cinfo;
   struct jpeg_error_mgr jerr;
@@ -384,15 +399,24 @@ int libjpeg_(Main_save)(lua_State *L) {
 
   /* this is a pointer to one row of image data */
   JSAMPROW row_pointer[1];
-  FILE *outfile = fopen( filename, "wb" );
-
-  if ( !outfile ) {
-    printf("Error opening output jpeg file %s\n!", filename );
-    return -1;
+  FILE *outfile = NULL;
+  if (save_to_file == 1) {
+    outfile = fopen( filename, "wb" );
+    if ( !outfile ) {
+      printf("Error opening output jpeg file %s\n!", filename );
+      return -1;
+    }
   }
+
   cinfo.err = jpeg_std_error( &jerr );
   jpeg_create_compress(&cinfo);
-  jpeg_stdio_dest(&cinfo, outfile);
+
+  /* specify data source (eg, a file) */
+  if (save_to_file == 1) {
+    jpeg_stdio_dest(&cinfo, outfile);
+  } else {
+    jpeg_mem_dest(&cinfo, &inmem, &inmem_size);
+  }
 
   /* Setting the parameters of the output file here */
   cinfo.image_width = width;	
@@ -402,6 +426,7 @@ int libjpeg_(Main_save)(lua_State *L) {
 
   /* default compression parameters, we shouldn't be worried about these */
   jpeg_set_defaults( &cinfo );
+  jpeg_set_quality(&cinfo, quality, (boolean)0);
 
   /* Now do the compression .. */
   jpeg_start_compress( &cinfo, TRUE );
@@ -415,7 +440,18 @@ int libjpeg_(Main_save)(lua_State *L) {
   /* similar to read file, clean up after we're done compressing */
   jpeg_finish_compress( &cinfo );
   jpeg_destroy_compress( &cinfo );
-  fclose( outfile );
+  
+  if (outfile != NULL) {
+    fclose( outfile );
+  }
+
+  if (save_to_file == 0) {
+    
+    THByteTensor_resize1d(tensor_dest, inmem_size);  /* will fail if it's not a Byte Tensor */ 
+    unsigned char* tensor_dest_data = THByteTensor_data(tensor_dest); 
+    memcpy(tensor_dest_data, inmem, inmem_size);
+    free(inmem);
+  }
 
   /* some cleanup */
   free(raw_image);
diff --git a/init.lua b/init.lua
index 50241e0..1d569b8 100755
--- a/init.lua
+++ b/init.lua
@@ -174,7 +174,9 @@ local function saveJPG(filename, tensor)
    local a = torch.Tensor():resize(tensor:size()):copy(tensor)
    a.image.saturate(a) -- bound btwn 0 and 1
    a:mul(MAXVAL)       -- remap to [0..255]
-   a.libjpeg.save(filename, a)
+   local save_to_file = 1
+   local quality = 75
+   a.libjpeg.save(filename, a, save_to_file, quality)
 end
 rawset(image, 'saveJPG', saveJPG)
 
@@ -185,6 +187,30 @@ function image.getJPGsize(filename)
    return torch.Tensor().libjpeg.size(filename)
 end
 
+local function compressJPG(tensor, quality)
+   if not xlua.require 'libjpeg' then
+      dok.error('libjpeg package not found, please install libjpeg',
+         'image.compressJPG')
+   end
+   local MAXVAL = 255
+   local a = torch.Tensor():resize(tensor:size()):copy(tensor)
+   a.image.saturate(a) -- bound btwn 0 and 1
+   a:mul(MAXVAL)       -- remap to [0..255]
+   local b = torch.ByteTensor()
+   local save_to_file = 0
+   quality = quality or 75
+   a.libjpeg.save("", a, save_to_file, quality, b)
+   return b
+end
+rawset(image, 'compressJPG', compressJPG)
+
+function image.getJPGsize(filename)
+   if not xlua.require 'libjpeg' then
+      dok.error('libjpeg package not found, please install libjpeg','image.getJPGsize')
+   end
+   return torch.Tensor().libjpeg.size(filename)
+end
+
 local function loadPPM(filename, depth, tensortype)
    require 'libppm'
    local MAXVAL = 255
diff --git a/test/test_compress_jpg.lua b/test/test_compress_jpg.lua
new file mode 100755
index 0000000..8a64bd8
--- /dev/null
+++ b/test/test_compress_jpg.lua
@@ -0,0 +1,43 @@
+require 'image'
+require 'paths'
+
+torch.setdefaulttensortype('torch.DoubleTensor')
+torch.setnumthreads(4)
+
+-- Create an instance of the test framework
+local mytester = torch.Tester()
+local precision_mean = 1e-3
+local precision_std = 1e-1
+local test = {}
+
+function test.CompressAndDecompress()
+  -- This test is unfortunately a correlated test: it will only be valid
+  -- if decompressJPG is OK.  However, since decompressJPG has it's own unit
+  -- test, this is problably fine. 
+  
+  local img = image.lena()
+
+  local quality = 100
+  local img_compressed = image.compressJPG(img, quality)
+  local size_100 = img_compressed:size(1)
+  local img_decompressed = image.decompressJPG(img_compressed)
+  local err = img_decompressed - img 
+
+  -- Now in general we will get BIG compression artifacts (even at quality=100)
+  -- but they will be relatively small, so instead of a abs():max() test, we do
+  -- a mean and std test.
+  local mean_err = err:mean()
+  local std_err = err:std()
+  mytester:assertlt(mean_err, precision_mean, 'compressJPG error is too high! ')
+  mytester:assertlt(std_err, precision_std, 'compressJPG error is too high! ')
+
+  -- Also check that the quality setting scales the size of the compressed image
+  quality = 25
+  img_compressed = image.compressJPG(img, quality)
+  local size_25 = img_compressed:size(1)
+  mytester:assertlt(size_25, size_100, 'compressJPG quality setting error! ')
+end
+
+-- Now run the test above
+mytester:add(test)
+mytester:run()