Added user specified pad value to warp function.

4 files changed, 28 insertions, 8 deletions

diff --git a/README.md b/README.md
index 3912801..df75e78 100644
--- a/README.md
+++ b/README.md
@@ -174,7 +174,7 @@ images requiring parameter Tensors like a warp `field` or a convolution
 `kernel`.
 
 <a name="image.warp"/>
-### [res] image.warp([dst,]src,field,[mode,offset,clamp]) ###
+### [res] image.warp([dst,]src,field,[mode,offset,clamp_mode,pad_val]) ###
 Warps image `src` (of size`KxHxW`) 
 according to flow field `field`. The latter has size `2xHxW` where the 
 first dimension is for the `(y,x)` flow field. String `mode` can 
@@ -182,8 +182,9 @@ take on values [lanczos](https://en.wikipedia.org/wiki/Lanczos_resampling),
 [bicubic](https://en.wikipedia.org/wiki/Bicubic_interpolation),
 [bilinear](https://en.wikipedia.org/wiki/Bilinear_interpolation) (the default), 
 or *simple*. When `offset` is true (the default), `(x,y)` is added to the flow field.
-The `clamp` variable specifies how to handle the interpolation of samples off the input image.
-Permitted values are strings *clamp* (the default) or *pad*. 
+The `clamp_mode` variable specifies how to handle the interpolation of samples off the input image.
+Permitted values are strings *clamp* (the default) or *pad*.
+When `clamp_mode` equals `pad`, the user can specify the padding value with `pad_val` (default = 0). Note: setting this value when `clamp_mode` equals `clamp` will result in an error.
 If `dst` is specified, it is used to store the result of the warp.
 Otherwise, returns a new `res` Tensor.
 
diff --git a/generic/image.c b/generic/image.c
index b06f664..6b69900 100755
--- a/generic/image.c
+++ b/generic/image.c
@@ -1644,6 +1644,7 @@ int image_(Main_warp)(lua_State *L) {
   int mode = lua_tointeger(L, 4);
   int offset_mode = lua_toboolean(L, 5);
   int clamp_mode = lua_tointeger(L, 6);
+  real pad_value = (real)lua_tonumber(L, 7);
 
   // dims
   int width = dst->size[2];
@@ -1680,7 +1681,7 @@ int image_(Main_warp)(lua_State *L) {
       if (off_image == 1 && clamp_mode == 1) {
         // We're off the image and we're clamping the input image to 0
         for (k=0; k<channels; k++) {
-          dst_data[ k*os[0] + y*os[1] + x*os[2] ] = 0;
+          dst_data[ k*os[0] + y*os[1] + x*os[2] ] = pad_value;
         }
       } else {
         ix = MAX(ix,0); ix = MIN(ix,src_width-1);
diff --git a/init.lua b/init.lua
index 7f32f98..8ea13f9 100644
--- a/init.lua
+++ b/init.lua
@@ -766,6 +766,7 @@ local function warp(...)
    local mode = 'bilinear'
    local offset_mode = true
    local clamp_mode = 'clamp'
+   local pad_value = 0
    local args = {...}
    local nargs = select('#',...)
    local bad_args = false
@@ -780,7 +781,12 @@ local function warp(...)
          mode = args[3]
          if nargs >= 4 then offset_mode = args[4] end
          if nargs >= 5 then clamp_mode = args[5] end
-         if nargs >= 6 then bad_args = true end
+         if nargs >= 6 then
+           assert(clamp_mode == 'pad', 'pad_value can only be specified if' ..
+                                       ' clamp_mode = "pad"')
+           pad_value = args[6]
+         end
+         if nargs >= 7 then bad_args = true end
       else
          -- With Destination tensor
          dst = args[1]
@@ -789,7 +795,12 @@ local function warp(...)
          if nargs >= 4 then mode = args[4] end
          if nargs >= 5 then offset_mode = args[5] end
          if nargs >= 6 then clamp_mode = args[6] end
-         if nargs >= 7 then bad_args = true end
+         if nargs >= 7 then
+           assert(clamp_mode == 'pad', 'pad_value can only be specified if' ..
+                                       ' clamp_mode = "pad"')
+           pad_value = args[7]
+         end
+         if nargs >= 8 then bad_args = true end
       end
    end
    if bad_args then
@@ -806,7 +817,8 @@ local function warp(...)
          {type='torch.Tensor', help='(y,x) flow field (2xHxW)', req=true},
          {type='string', help='mode: lanczos | bicubic | bilinear | simple', default='bilinear'},
          {type='string', help='offset mode (add (x,y) to flow field)', default=true},
-         {type='string', help='clamp mode: how to handle interp of samples off the input image (clamp | pad)', default='clamp'}))
+         {type='string', help='clamp mode: how to handle interp of samples off the input image (clamp | pad)', default='clamp'},
+         {type='number', help='pad value: value to pad image. Can only be set when clamp mode equals "pad"', default=0}))
       dok.error('incorrect arguments', 'image.warp')
    end
    -- This is a little messy, but convert mode string to an enum
@@ -837,7 +849,7 @@ local function warp(...)
    dst = dst or src.new()
    dst:resize(src:size(1), field:size(2), field:size(3))
 
-   src.image.warp(dst, src, field, mode, offset_mode, clamp_mode)
+   src.image.warp(dst, src, field, mode, offset_mode, clamp_mode, pad_value)
    if dim2 then
       dst = dst[1]
    end
diff --git a/test/test_warp.lua b/test/test_warp.lua
index 68331c8..68a4ac1 100644
--- a/test/test_warp.lua
+++ b/test/test_warp.lua
@@ -129,5 +129,11 @@ t1 = sys.clock()
 print("Rotation Time lanczos = " .. (t1 - t0))  -- Not a robust measure (should average)
 image.display{image = im_lanczos, zoom = 4, legend = 'rotation lanczos'}
 
+im_lanczos = image.warp(im, flow, 'lanczos', false, 'pad')
+image.display{image = im_lanczos, zoom = 4, legend = 'rotation lanczos (default pad)'}
+
+im_lanczos = image.warp(im, flow, 'lanczos', false, 'pad', 1)
+image.display{image = im_lanczos, zoom = 4, legend = 'rotation lanczos (pad 1)'}
+
 image.display{image = im, zoom = 4, legend = 'source image'}