added 3-channels support to cv::setIdentity

mjul · Mar 18, 2014 · 04884eb · 04884eb
1 parent 8d97d0d
commit 04884eb
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 9 deletions.
diff --git a/modules/core/src/matrix.cpp b/modules/core/src/matrix.cpp
@@ -2679,17 +2679,17 @@ namespace cv {
 
 static bool ocl_setIdentity( InputOutputArray _m, const Scalar& s )
 {
-    int type = _m.type(), cn = CV_MAT_CN(type);
-    if (cn == 3)
-        return false;
+    int type = _m.type(), depth = CV_MAT_DEPTH(type), cn = CV_MAT_CN(type),
+            sctype = CV_MAKE_TYPE(depth, cn == 3 ? 4 : cn);
 
     ocl::Kernel k("setIdentity", ocl::core::set_identity_oclsrc,
-                  format("-D T=%s", ocl::memopTypeToStr(type)));
+                  format("-D T=%s -D T1=%s -D cn=%d -D ST=%s", ocl::memopTypeToStr(type),
+                         ocl::memopTypeToStr(depth), cn, ocl::memopTypeToStr(sctype)));
     if (k.empty())
         return false;
 
     UMat m = _m.getUMat();
-    k.args(ocl::KernelArg::WriteOnly(m), ocl::KernelArg::Constant(Mat(1, 1, type, s)));
+    k.args(ocl::KernelArg::WriteOnly(m), ocl::KernelArg::Constant(Mat(1, 1, sctype, s)));
 
     size_t globalsize[2] = { m.cols, m.rows };
     return k.run(2, globalsize, NULL, false);

diff --git a/modules/core/src/opencl/set_identity.cl b/modules/core/src/opencl/set_identity.cl
@@ -43,17 +43,28 @@
 //
 //M*/
 
+#if cn != 3
+#define loadpix(addr) *(__global const T *)(addr)
+#define storepix(val, addr)  *(__global T *)(addr) = val
+#define TSIZE (int)sizeof(T)
+#define scalar scalar_
+#else
+#define loadpix(addr) vload3(0, (__global const T1 *)(addr))
+#define storepix(val, addr) vstore3(val, 0, (__global T1 *)(addr))
+#define TSIZE ((int)sizeof(T1)*3)
+#define scalar (T)(scalar_.x, scalar_.y, scalar_.z)
+#endif
+
 __kernel void setIdentity(__global uchar * srcptr, int src_step, int src_offset, int rows, int cols,
-                          T scalar)
+                          ST scalar_)
 {
     int x = get_global_id(0);
     int y = get_global_id(1);
 
     if (x < cols && y < rows)
     {
-        int src_index = mad24(y, src_step, mad24(x, (int)sizeof(T), src_offset));
-        __global T * src = (__global T *)(srcptr + src_index);
+        int src_index = mad24(y, src_step, mad24(x, TSIZE, src_offset));
 
-        src[0] = x == y ? scalar : (T)(0);
+        storepix(x == y ? scalar : (T)(0), srcptr + src_index);
     }
 }