Merge pull request #60 from jonathantompson/flip

Added flip function (to a tensor up to 5D along an arbitrary dimension).

4 files changed, 224 insertions, 0 deletions

diff --git a/README.md b/README.md
index 314062c..9c71d7e 100644..100755
--- a/README.md
+++ b/README.md
@@ -128,6 +128,11 @@ store the output image. Otherwise, returns a new `res` Tensor.
 Flips image `src` vertically (upsize<->down). If `dst` is provided, it is used to
 store the output image. Otherwise, returns a new `res` Tensor.
 
+<a name="image.flip"/>
+### [res] image.flip([dst,] src, flip_dim) ###
+Flips image `src` along the specified dimension. If `dst` is provided, it is used to
+store the output image. Otherwise, returns a new `res` Tensor.
+
 <a name="image.minmax"/>
 ### [res] image.minmax{tensor, [min, max, ...]} ###
 Compresses image `tensor` between `min` and `max`. 
diff --git a/generic/image.c b/generic/image.c
index dfccfa2..2faab8c 100644..100755
--- a/generic/image.c
+++ b/generic/image.c
@@ -1425,6 +1425,79 @@ int image_(Main_hflip)(lua_State *L) {
   return 0;
 }
 
+/* flip an image along a specified dimension */
+int image_(Main_flip)(lua_State *L) {
+  THTensor *dst = luaT_checkudata(L, 1, torch_Tensor);
+  THTensor *src = luaT_checkudata(L, 2, torch_Tensor);
+  long flip_dim = luaL_checklong(L, 3);
+  
+  if (dst->nDimension != src->nDimension) {
+    luaL_error(L, "image.flip: src and dst nDimension does not match");
+  }
+  
+  if (flip_dim < 1 || flip_dim > dst->nDimension) {
+    luaL_error(L, "image.flip: flip_dim out of bounds");
+  }
+  flip_dim--;  //  Make it zero indexed
+
+  // get raw pointers
+  real *dst_data = THTensor_(data)(dst);
+  real *src_data = THTensor_(data)(src);
+  if (dst_data == src_data) {  
+    luaL_error(L, "image.flip: in-place flip not supported");
+  }
+  
+  long size0 = dst->size[0];
+  long size1 = dst->size[1];
+  long size2 = dst->size[2];
+  long size3 = dst->size[3];
+  long size4 = dst->size[4];
+  long size_flip = dst->size[flip_dim];
+  
+  if (src->size[0] != size0 || src->size[1] != size1 || 
+      src->size[2] != size2 || src->size[3] != size3 || 
+      src->size[4] != size4) {
+    luaL_error(L, "image.flip: src and dst are not the same size");
+  }
+  
+  long *is = src->stride;
+  long *os = dst->stride;
+
+  long x, y, z, d, t, isrc, idst;
+  for (t = 0; t < size0; t++) {
+    for (d = 0; d < size1; d++) {
+      for (z = 0; z < size2; z++) {
+        for (y = 0; y < size3; y++) {
+          for (x = 0; x < size4; x++) {
+            isrc = t*is[0] + d*is[1] + z*is[2] + y*is[3] + x*is[4];
+            // The big switch statement here looks ugly, however on my machine
+            // gcc compiles it to a skip list, so it should be fast.
+            switch (flip_dim) {
+              case 0:
+                idst = (size0 - t - 1)*os[0] + d*os[1] + z*os[2] + y*os[3] + x*os[4];
+                break;
+              case 1:
+                idst = t*os[0] + (size1 - d - 1)*os[1] + z*os[2] + y*os[3] + x*os[4];
+                break;
+              case 2:
+                idst = t*os[0] + d*os[1] + (size2 - z - 1)*os[2] + y*os[3] + x*os[4];
+                break;
+              case 3:
+                idst = t*os[0] + d*os[1] + z*os[2] + (size3 - y - 1)*os[3] + x*os[4];
+                break;
+              case 4:
+                idst = t*os[0] + d*os[1] + z*os[2] + y*os[3] + (size4 - x - 1)*os[4];
+                break;
+            }              
+            dst_data[ idst ] = src_data[  isrc ];
+          }
+        }
+      }
+    }
+  }
+
+  return 0;
+}
 
 /*
  * Warps an image, according to an (x,y) flow field. The flow
@@ -1787,6 +1860,7 @@ static const struct luaL_Reg image_(Main__) [] = {
   {"gaussian", image_(Main_gaussian)},
   {"vflip", image_(Main_vflip)},
   {"hflip", image_(Main_hflip)},
+  {"flip", image_(Main_flip)},
   {"colorize", image_(Main_colorize)},
   {NULL, NULL}
 };
diff --git a/init.lua b/init.lua
index a7f48a1..2bf9fb3 100755
--- a/init.lua
+++ b/init.lua
@@ -811,12 +811,22 @@ local function hflip(...)
                        {type='torch.Tensor', help='input image', req=true}))
       dok.error('incorrect arguments', 'image.hflip')
    end
+
+   if not src:isContiguous() then
+     dok.error('input tensor is not contiguous', 'image.hflip')
+   end
+
    dst = dst or src.new()
    local original_size = src:size()
    if src:nDimension() == 2 then
       src = src:new():resize(1,src:size(1),src:size(2))
    end
    dst:resizeAs(src)
+
+   if not dst:isContiguous() then
+     dok.error('destination tensor is not contiguous', 'image.hflip')
+   end
+
    dst.image.hflip(dst, src)
    dst:resize(original_size)
    return dst
@@ -843,12 +853,22 @@ local function vflip(...)
                        {type='torch.Tensor', help='input image', req=true}))
       dok.error('incorrect arguments', 'image.vflip')
    end
+
+   if not src:isContiguous() then
+     dok.error('input tensor is not contiguous', 'image.vflip')
+   end
+
    dst = dst or src.new()
    local original_size = src:size()
    if src:nDimension() == 2 then
       src = src:new():resize(1,src:size(1),src:size(2))
    end
    dst:resizeAs(src)
+  
+   if not dst:isContiguous() then
+     dok.error('destination tensor is not contiguous', 'image.vflip')
+   end
+ 
    dst.image.vflip(dst, src)
    dst:resize(original_size)
    return dst
@@ -856,6 +876,69 @@ end
 rawset(image, 'vflip', vflip)
 
 ----------------------------------------------------------------------
+-- flip (specify dimension, up to 5D tensor)
+--
+local function flip(...)
+   local dst,src,flip_dim
+   local args = {...}
+   if select('#',...) == 3 then
+      dst = args[1]
+      src = args[2]
+      flip_dim = args[3]
+   elseif select('#',...) == 2 then
+      src = args[1]
+      flip_dim = args[2]
+   else
+      print(dok.usage('image.flip',
+                       'flips an image along a specified dimension', nil,
+                       {type='torch.Tensor', help='input image', req=true},
+                       {type='number', help='Dimension to flip', req=true},
+                       '',
+                       {type='torch.Tensor', help='destination', req=true},
+                       {type='torch.Tensor', help='input image', req=true},
+                       {type='number', help='Dimension to flip', req=true}))
+      dok.error('incorrect arguments', 'image.flip')
+   end
+   assert(src:nDimension() <= 5, 'too many input dims (up to 5D supported)')
+   assert(flip_dim <= src:nDimension() and flip_dim >= 1, 'Bad flip dimension')
+
+   if not src:isContiguous() then
+     dok.error('input tensor is not contiguous', 'image.flip')
+   end
+
+   dst = dst or src.new()
+   local original_size = src:size()
+   local flip_dim_cpp
+   if src:nDimension() == 1 then
+      src = src:new():resize(1, 1, 1, 1, src:size(1))
+      flip_dim_cpp = flip_dim + 4
+   elseif src:nDimension() == 2 then
+      src = src:new():resize(1, 1, 1, src:size(1), src:size(2))
+      flip_dim_cpp = flip_dim + 3
+   elseif src:nDimension() == 3 then
+      src = src:new():resize(1, 1, src:size(1), src:size(2),src:size(3))
+      flip_dim_cpp = flip_dim + 2
+   elseif src:nDimension() == 4 then
+      src = src:new():resize(1, src:size(1), src:size(2), src:size(3), 
+        src:size(4))
+      flip_dim_cpp = flip_dim + 1
+   else
+      flip_dim_cpp = flip_dim
+   end
+   dst:resizeAs(src)
+
+   if not dst:isContiguous() then
+     dok.error('destination tensor is not contiguous', 'image.flip')
+   end
+
+   dst.image.flip(dst, src, flip_dim_cpp)
+   dst:resize(original_size)
+   return dst
+end
+
+rawset(image, 'flip', flip)
+
+----------------------------------------------------------------------
 -- convolve(dst,src,ker,type)
 -- convolve(dst,src,ker)
 -- dst = convolve(src,ker,type)
diff --git a/test/test_flip.lua b/test/test_flip.lua
new file mode 100755
index 0000000..4f9476d
--- /dev/null
+++ b/test/test_flip.lua
@@ -0,0 +1,62 @@
+require 'image'
+
+torch.setdefaulttensortype('torch.DoubleTensor')
+torch.setnumthreads(8)
+
+-- Create an instance of the test framework
+local precision = 1e-5
+local mytester = torch.Tester()
+local test = {}
+
+-- This is a correlated test (which is kinda lazy).  We're assuming HFLIP is OK.
+function test.FlipAgainstHFlip()
+  for ndims = 1, 5 do
+    for flip_dim = 1, ndims do
+      local sz = {}
+      for i = 1, ndims do
+        sz[i] = math.random(5,10)
+      end
+      
+      local input = torch.rand(unpack(sz))
+      local output = image.flip(input, flip_dim)
+      
+      -- Now perform the same operation using HFLIP
+      local input_tran = input
+      if (flip_dim < ndims) then
+        -- First permute the flip dimension to X dim
+        input_tran = input:transpose(flip_dim, ndims):contiguous()
+      end
+      -- Now reshape it to 3D
+      local original_hflip_sz = input_tran:size()
+      if ndims == 1 then
+        input_tran:resize(1, original_hflip_sz[1])
+      end
+      if ndims > 3 then
+        sz1 = 1
+        for i = 1, ndims - 2 do
+          sz1 = sz1 * original_hflip_sz[i]
+        end
+        input_tran:resize(sz1, original_hflip_sz[input_tran:dim()-1], 
+          original_hflip_sz[input_tran:dim()])
+      end
+     
+      local output_hflip = image.hflip(input_tran)
+      
+      -- Put it back to Ndim
+      output_hflip:resize(original_hflip_sz)
+      
+      if (flip_dim < ndims) then
+        -- permute bacx the flip dimension
+        output_hflip = output_hflip:transpose(flip_dim, ndims):contiguous()
+      end
+      
+      local err = output_hflip - output
+      mytester:asserteq(err:abs():max(), 0, 'error - bad flip! (ndims='..
+        ndims..',flip_dim='..flip_dim..')')
+    end
+  end
+end
+
+-- Now run the test above
+mytester:add(test)
+mytester:run()