rgb2y for byte tensors (needed to go into C for efficiency)bytechan1

2 files changed, 34 insertions, 2 deletions

diff --git a/generic/image.c b/generic/image.c
index 2faab8c..8aa77f0 100755
--- a/generic/image.c
+++ b/generic/image.c
@@ -1838,6 +1838,37 @@ int image_(Main_colorize)(lua_State *L) {
   return 0;
 }
 
+int image_(Main_rgb2y)(lua_State *L) {
+  THTensor *rgb = luaT_checkudata(L, 1, torch_Tensor);
+  THTensor *yim = luaT_checkudata(L, 2, torch_Tensor);
+
+  luaL_argcheck(L, rgb->nDimension == 3, 1, "image.rgb2y: src not 3D");
+  luaL_argcheck(L, yim->nDimension == 2, 2, "image.rgb2y: dst not 2D");
+  luaL_argcheck(L, rgb->size[1] == yim->size[0], 2,
+		"image.rgb2y: src and dst not of same height");
+  luaL_argcheck(L, rgb->size[2] == yim->size[1], 2,
+		"image.rgb2y: src and dst not of same width");
+
+  int y,x;
+  real r,g,b,yc;
+  const int height = rgb->size[1];
+  const int width = rgb->size[2];
+  for (y=0; y<height; y++) {
+    for (x=0; x<width; x++) {
+      // get Rgb
+      r = THTensor_(get3d)(rgb, 0, y, x);
+      g = THTensor_(get3d)(rgb, 1, y, x);
+      b = THTensor_(get3d)(rgb, 2, y, x);
+
+      yc = (real) ((0.299 * (float) r)
+		   + (0.587 * (float) g)
+		   + (0.114 * (float) b));
+      THTensor_(set2d)(yim, y, x, yc);
+    }
+  }
+  return 0;
+}
+
 static const struct luaL_Reg image_(Main__) [] = {
   {"scaleSimple", image_(Main_scaleSimple)},
   {"scaleBilinear", image_(Main_scaleBilinear)},
@@ -1851,6 +1882,7 @@ static const struct luaL_Reg image_(Main__) [] = {
   {"cropNoScale", image_(Main_cropNoScale)},
   {"warp", image_(Main_warp)},
   {"saturate", image_(Main_saturate)},
+  {"rgb2y",   image_(Main_rgb2y)},
   {"rgb2hsv", image_(Main_rgb2hsv)},
   {"rgb2hsl", image_(Main_rgb2hsl)},
   {"hsv2rgb", image_(Main_hsv2rgb)},
diff --git a/init.lua b/init.lua
index 6c54281..330aa56 100755
--- a/init.lua
+++ b/init.lua
@@ -61,7 +61,7 @@ local function todepth(img, depth)
       if img:nDimension() == 2 then
          -- all good
       elseif img:size(1) == 3 or img:size(1) == 4 then
-         img = image.rgb2y(img:narrow(1,1,3))[1]
+	 img = image.rgb2y(img:narrow(1,1,3))[1]
       elseif img:size(1) == 2 then
          img = img:narrow(1,1,1)
       elseif img:size(1) ~= 1 then
@@ -1468,7 +1468,7 @@ function image.rgb2y(...)
    local outputY = output[1]
 
    -- convert
-   outputY:zero():add(0.299, inputRed):add(0.587, inputGreen):add(0.114, inputBlue)
+   input.image.rgb2y(input, outputY)
 
    -- return YUV image
    return output