Implement drawRect function

7 files changed, 103 insertions, 2 deletions

diff --git a/README.md b/README.md
index cab3821..97d505b 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,7 @@ images. It contains a wide variety of functions divided into the following categ
   * [Saving and loading](doc/saveload.md) images as JPEG, PNG, PPM and PGM;
   * [Simple transformations](doc/simpletransform.md) like translation, scaling and rotation;
   * [Parameterized transformations](doc/paramtransform.md) like convolutions and warping;
-  * [Simple Drawing Routines](doc/drawing.md) like drawing text on an image;
+  * [Simple Drawing Routines](doc/drawing.md) like drawing text or a rectangle on an image;
   * [Graphical user interfaces](doc/gui.md) like display and window;
   * [Color Space Conversions](doc/colorspace.md) from and to RGB, YUV, Lab, and HSL;
   * [Tensor Constructors](doc/tensorconstruct.md) for creating Lenna, Fabio and Gaussian and Laplacian kernels;
diff --git a/assets/rectangle.png b/assets/rectangle.png
new file mode 100644
index 0000000..aa4720e
--- /dev/null
+++ b/assets/rectangle.png
diff --git a/doc/drawing.md b/doc/drawing.md
index de1044c..90aca0b 100644
--- a/doc/drawing.md
+++ b/doc/drawing.md
@@ -22,3 +22,23 @@ image.drawText(image.lena(), "hello\nworld", 10, 10)
 image.drawText(image.lena(), "hello\nworld", 10, 20,{color = {0, 255, 0}, size = 5})
 image.drawText(image.lena(), "hello\nworld", 10, 20,{color = {0, 255, 0}, bg = {255, 0, 0}, size = 5})
 ```
+
+<a name="image.drawRect"></a>
+### [res] image.drawRect(src, x1, y1, x2, y2, [options]) ###
+Draws a rectangle onto a 3-channel Tensor (C x H x W).  The top-left corner of
+the rectangle is `x1, y1`, and the bottom-right corner is `x2, y2`.
+
+The `options` table can be passed in to set color, in-place etc.
+
+Options:
+* `color` - [table] The rectangle color. A table of 3 numbers `{R, G, B}`, each
+    number scaled between 0 and 255. For example, `red` is `{255, 0, 0}`
+* `lineWidth` - [number] The width of the rectangle line, in pixels
+* `inplace` - [boolean] If true, draws directly on the input tensor and returns
+    it. `default value = false`
+
+Example:
+
+```lua
+image.drawRect(image.lena(), 200, 200, 370, 400, {lineWidth = 5, color = {0, 255, 0}})
+```
diff --git a/doc/index.md b/doc/index.md
index 62a987d..5e50c39 100644
--- a/doc/index.md
+++ b/doc/index.md
@@ -6,7 +6,7 @@ images. It contains a wide variety of functions divided into the following categ
   * [Saving and loading](saveload.md) images as JPEG, PNG, PPM and PGM;
   * [Simple transformations](simpletransform.md) like translation, scaling and rotation;
   * [Parameterized transformations](paramtransform.md) like convolutions and warping;
-  * [Simple Drawing Routines](doc/drawing.md) like drawing text on an image;
+  * [Simple Drawing Routines](doc/drawing.md) like drawing text or a rectangle on an image;
   * [Graphical user interfaces](gui.md) like display and window;
   * [Color Space Conversions](colorspace.md) from and to RGB, YUV, Lab, and HSL;
   * [Tensor Constructors](tensorconstruct.md) for creating Lenna, Fabio and Gaussian and Laplacian kernels;
diff --git a/generic/image.c b/generic/image.c
index 0990814..f30fcad 100755
--- a/generic/image.c
+++ b/generic/image.c
@@ -2216,6 +2216,45 @@ int image_(Main_drawtext)(lua_State *L) {
   return 0;
 }
 
+int image_(Main_drawRect)(lua_State *L) {
+  THTensor *output = (THTensor *)luaT_checkudata(L, 1, torch_Tensor);
+  long x1long = luaL_checklong(L, 2);
+  long y1long = luaL_checklong(L, 3);
+  long x2long = luaL_checklong(L, 4);
+  long y2long = luaL_checklong(L, 5);
+  int lineWidth = luaL_checkint(L, 6);
+  int cr = luaL_checkint(L, 7);
+  int cg = luaL_checkint(L, 8);
+  int cb = luaL_checkint(L, 9);
+
+  int offset = lineWidth / 2;
+  int x1 = (int) MAX(0, x1long - offset - 1);
+  int y1 = (int) MAX(0, y1long - offset - 1);
+  int x2 = (int) MIN(output->size[2] - 1, x2long - offset - 1);
+  int y2 = (int) MIN(output->size[1] - 1, y2long - offset - 1);
+
+  int w = x2 - x1 + 1;
+  int h = y2 - y1 + 1;
+  for (int y = y1; y < y2 + lineWidth; y++) {
+    for (int x = x1; x < x1 + lineWidth; x++) {
+      image_(drawPixel)(output, y, x, cr, cg, cb);
+    }
+    for (int x = x2; x < x2 + lineWidth; x++) {
+      image_(drawPixel)(output, y, x, cr, cg, cb);
+    }
+  }
+  for (int x = x1; x < x2 + lineWidth; x++) {
+    for (int y = y1; y < y1 + lineWidth; y++) {
+      image_(drawPixel)(output, y, x, cr, cg, cb);
+    }
+    for (int y = y2; y < y2 + lineWidth; y++) {
+      image_(drawPixel)(output, y, x, cr, cg, cb);
+    }
+  }
+
+  return 0;
+}
+
 
 static const struct luaL_Reg image_(Main__) [] = {
   {"scaleSimple", image_(Main_scaleSimple)},
@@ -2244,6 +2283,7 @@ static const struct luaL_Reg image_(Main__) [] = {
   {"flip", image_(Main_flip)},
   {"colorize", image_(Main_colorize)},
   {"text", image_(Main_drawtext)},
+  {"drawRect", image_(Main_drawRect)},
   {NULL, NULL}
 };
 
diff --git a/init.lua b/init.lua
index 80b5f06..8a14cec 100644
--- a/init.lua
+++ b/init.lua
@@ -1977,6 +1977,27 @@ function image.drawText(src, text, x, y, opts)
 end
 
 ----------------------------------------------------------------------
+--- Draw a rectangle on the image
+--
+-- color, bgcolor, size, wrap, inplace
+function image.drawRect(src, x1, y1, x2, y2, opts)
+   opts = opts or {}
+   assert(torch.isTensor(src) and src:dim() == 3 and src:size(1) == 3,
+    "input image has to be a 3D tensor of shape 3 x H x W ")
+   local out = src
+   if not opts.inplace then
+      out = src:clone()
+   end
+   if not (x1 and x2 and y1 and y2) then return out end
+   local color = opts.color or {255, 0, 0} -- red default
+   local lineWidth = opts.lineWidth or 1
+
+   src.image.drawRect(out, x1, y1, x2, y2, lineWidth, color[1], color[2], color[3])
+   return out
+end
+
+
+----------------------------------------------------------------------
 --- Returns a gaussian kernel.
 --
 function image.gaussian(...)
diff --git a/test/test.lua b/test/test.lua
index 8040673..a5ca949 100644
--- a/test/test.lua
+++ b/test/test.lua
@@ -652,6 +652,26 @@ function test.test_textdraw()
   end
 end
 
+----------------------------------------------------------------------
+-- Text drawing rect
+--
+function test.test_drawRect()
+  local types = {
+     ["torch.ByteTensor"]   = "byte",
+     ["torch.DoubleTensor"] = "double",
+     ["torch.FloatTensor"]  = "float"
+  }
+  for k,v in pairs(types) do
+    local bg = torch.zeros(3, 24, 12):type(k)
+    if k == 'torch.ByteTensor' then
+      bg:fill(3)
+    else
+      bg:fill(3/255)
+    end
+    local img = image.drawRect(bg, 5, 5, 10, 20, {color={255, 0, 255}})
+    checkPNG(getTestImagePath("rectangle.png"), 3, v, img)
+  end
+end
 
 function image.test(tests, seed)
    local defaultTensorType = torch.getdefaulttensortype()