Add image.getSize(filename) function

Add getPPMsize and getPGMsize functions
Combine them with getPNGsize and getJPGsize for getSize function

3 files changed, 121 insertions, 11 deletions

diff --git a/doc/saveload.md b/doc/saveload.md
index d90fc20..b9fa1fd 100644
--- a/doc/saveload.md
+++ b/doc/saveload.md
@@ -1,23 +1,23 @@
 <a name="image.saveload"></a>
 ## Saving and Loading ##
-This sections includes functions for saving and loading different types 
+This sections includes functions for saving and loading different types
 of images to and from disk.
 
 <a name="image.load"></a>
 ### [res] image.load(filename, [depth, tensortype]) ###
 Loads an image located at path `filename` 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 
+of type `tensortype` (*float*, *double* or *byte*). The last two arguments
 are optional.
 
-The image format is determined from the `filename`'s 
-extension suffix. Supported formats are 
-[JPEG](https://en.wikipedia.org/wiki/JPEG), 
-[PNG](https://en.wikipedia.org/wiki/Portable_Network_Graphics), 
+The image format is determined from the `filename`'s
+extension suffix. Supported formats are
+[JPEG](https://en.wikipedia.org/wiki/JPEG),
+[PNG](https://en.wikipedia.org/wiki/Portable_Network_Graphics),
 [PPM and PGM](https://en.wikipedia.org/wiki/Netpbm_format).
- 
-The returned `res` Tensor has size `nChannel x height x width` where `nChannel` is 
-1 (greyscale) or 3 (usually [RGB](https://en.wikipedia.org/wiki/RGB_color_model) 
+
+The returned `res` Tensor has size `nChannel x height x width` where `nChannel` is
+1 (greyscale) or 3 (usually [RGB](https://en.wikipedia.org/wiki/RGB_color_model)
 or [YUV](https://en.wikipedia.org/wiki/YUV).
 
 Usage:
@@ -30,9 +30,21 @@ local img = image.load(imagefile,1,'byte')
 
 ```
 
+<a name="image.getSize"></a>
+### [res] image.getSize(filename) ###
+Return the size of an image located at path `filename` into a LongTensor.
+
+The image format is determined from the `filename`'s
+extension suffix. Supported formats are
+[JPEG](https://en.wikipedia.org/wiki/JPEG),
+[PNG](https://en.wikipedia.org/wiki/Portable_Network_Graphics),
+[PPM and PGM](https://en.wikipedia.org/wiki/Netpbm_format).
+
+The returned `res` Tensor has size `3` (nChannel, height, width).
+
 <a name="image.save"></a>
 ### image.save(filename, tensor) ###
-Saves Tensor `tensor` to disk at path `filename`. The format to which 
+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`.
 To save with a minimal loss, the tensor values should lie in the range [0, 1] since the tensor is clamped between 0 and 1 before being saved to the disk.
diff --git a/generic/ppm.c b/generic/ppm.c
index 6d324e9..7cd36e6 100644
--- a/generic/ppm.c
+++ b/generic/ppm.c
@@ -2,6 +2,50 @@
 #define TH_GENERIC_FILE "generic/ppm.c"
 #else
 
+static int libppm_(Main_size)(lua_State *L)
+{
+  const char *filename = luaL_checkstring(L, 1);
+  FILE* fp = fopen ( filename, "r" );
+  if ( !fp ) {
+    luaL_error(L, "cannot open file <%s> for reading", filename);
+  }
+
+  long W,H,C;
+  char p,n;
+
+  // magic number
+  p = (char)getc(fp);
+  if ( p != 'P' ) {
+    W = H = 0;
+    fclose(fp);
+    luaL_error(L, "corrupted file");
+  }
+
+  n = (char)getc(fp);
+
+  // Dimensions
+  W = ppm_get_long(fp);
+  H = ppm_get_long(fp);
+
+  if ( n=='3' || n == '6') {
+    C = 3;
+  } else if ( n=='2' || n=='5' ) {
+    C = 1;
+  } else {
+    W=H=C=0;
+    fclose ( fp );
+    luaL_error(L, "unsupported magic number: P%c", n);
+  }
+
+  fclose(fp);
+
+  lua_pushnumber(L, C);
+  lua_pushnumber(L, H);
+  lua_pushnumber(L, W);
+
+  return 3;
+}
+
 static int libppm_(Main_load)(lua_State *L)
 {
   const char *filename = luaL_checkstring(L, 1);
@@ -170,6 +214,7 @@ static const luaL_Reg libppm_(Main__)[] =
 {
   {"load", libppm_(Main_load)},
   {"save", libppm_(Main_save)},
+  {"size", libppm_(Main_size)},
   {NULL, NULL}
 };
 
diff --git a/init.lua b/init.lua
index 59c17a2..e77f309 100644
--- a/init.lua
+++ b/init.lua
@@ -309,6 +309,11 @@ local function savePGM(filename, tensor)
 end
 rawset(image, 'savePGM', savePGM)
 
+function image.getPPMsize(filename)
+   require 'libppm'
+   return torch.Tensor().libppm.size(filename)
+end
+
 local filetypes = {
    jpg = {loader = image.loadJPG, saver = image.saveJPG},
    png = {loader = image.loadPNG, saver = image.savePNG},
@@ -376,6 +381,54 @@ local function load(filename, depth, tensortype)
 end
 rawset(image, 'load', load)
 
+filetypes.jpg.sizer = image.getJPGsize
+filetypes.png.sizer = image.getPNGsize
+filetypes.ppm.sizer = image.getPPMsize
+filetypes.pgm.sizer = image.getPPMsize -- sim. to loadPPM not loadPGM
+
+local function getSize(filename)
+   if not filename then
+      print(dok.usage('image.getSize',
+                       'returns size of image without loading', nil,
+                       {type='string', help='path to file', req=true}))
+      dok.error('missing file name', 'image.getSize')
+   end
+
+   local ext
+
+   local f, err = io.open(filename, 'rb')
+   if not f then
+      error(err)
+   end
+   local hdr = f:read(4) or ''
+   f:close()
+
+   if startswith(hdr, magicJPG) then
+      ext = 'jpg'
+   elseif startswith(hdr, magicPNG) then
+      ext = 'png'
+   elseif hdr:match('^P[25]') then
+      ext = 'pgm'
+   elseif hdr:match('^P[36]') then
+      ext = 'ppm'
+   end
+
+   if not ext then
+      ext = string.match(filename,'%.(%a+)$')
+   end
+   local size
+   if image.is_supported(ext) then
+      size = {filetypes[ext].sizer(filename)}
+   elseif not ext then
+      dok.error('unable to determine image type for file: ' .. filename, 'image.getSize')
+   else
+      dok.error('unknown image type: ' .. ext, 'image.load')
+   end
+
+   return torch.LongTensor(size)
+end
+rawset(image, 'getSize', getSize)
+
 local function save(filename, tensor)
    if not filename or not tensor then
       print(dok.usage('image.save',
@@ -1076,7 +1129,7 @@ local function affinetransform(...)
    local matrix = matrix:typeAs(torch.Tensor()) -- coerce matrix to default tensor type
    local field = torch.mm(matrix, view_xy)
    field = (grid_xy - field:reshape( 2, height, width ))
-   
+
    -- offset field for translation
    translation = torch.Tensor(translation)
    field[1] = field[1] - translation[1]