added support of 3 channels images to StereoBeliefPropagation_GPU

diff --git a/modules/gpu/src/beliefpropagation_gpu.cpp b/modules/gpu/src/beliefpropagation_gpu.cpp
index 3841550..3ebf246 100644 (file)
--- a/modules/gpu/src/beliefpropagation_gpu.cpp
+++ b/modules/gpu/src/beliefpropagation_gpu.cpp
@@ -65,11 +65,11 @@ const float DEFAULT_DISC_SINGLE_JUMP = 1.0f;
 \r
 namespace cv { namespace gpu { namespace impl {\r
     void load_constants(int ndisp, float max_data_term, float data_weight, float max_disc_term, float disc_single_jump);\r
-    void comp_data(int msgType, const DevMem2D& l, const DevMem2D& r, DevMem2D mdata, const cudaStream_t& stream);\r
-    void data_step_down(int dst_cols, int dst_rows, int src_rows, int msgType, const DevMem2D& src, DevMem2D dst, const cudaStream_t& stream);\r
-    void level_up_messages(int dst_idx, int dst_cols, int dst_rows, int src_rows, int msgType, DevMem2D* mus, DevMem2D* mds, DevMem2D* mls, DevMem2D* mrs, const cudaStream_t& stream);\r
-    void calc_all_iterations(int cols, int rows, int iters, int msgType, DevMem2D& u, DevMem2D& d, DevMem2D& l, DevMem2D& r, const DevMem2D& data, const cudaStream_t& stream);\r
-    void output(int msgType, const DevMem2D& u, const DevMem2D& d, const DevMem2D& l, const DevMem2D& r, const DevMem2D& data, DevMem2D disp, const cudaStream_t& stream);\r
+    void comp_data(int msg_type, const DevMem2D& l, const DevMem2D& r, int channels, DevMem2D mdata, const cudaStream_t& stream);\r
+    void data_step_down(int dst_cols, int dst_rows, int src_rows, int msg_type, const DevMem2D& src, DevMem2D dst, const cudaStream_t& stream);\r
+    void level_up_messages(int dst_idx, int dst_cols, int dst_rows, int src_rows, int msg_type, DevMem2D* mus, DevMem2D* mds, DevMem2D* mls, DevMem2D* mrs, const cudaStream_t& stream);\r
+    void calc_all_iterations(int cols, int rows, int iters, int msg_type, DevMem2D& u, DevMem2D& d, DevMem2D& l, DevMem2D& r, const DevMem2D& data, const cudaStream_t& stream);\r
+    void output(int msg_type, const DevMem2D& u, const DevMem2D& d, const DevMem2D& l, const DevMem2D& r, const DevMem2D& data, DevMem2D disp, const cudaStream_t& stream);\r
 }}}\r
 \r
 cv::gpu::StereoBeliefPropagation_GPU::StereoBeliefPropagation_GPU(int ndisp_, int iters_, int levels_, int msg_type_, float msg_scale_)\r
@@ -228,7 +228,7 @@ static void stereo_bp_gpu_operator(int ndisp, int iters, int levels,
 \r
     datas[0].create(rows * ndisp, cols, msg_type);\r
 \r
-    impl::comp_data(msg_type, left, right, datas.front(), stream);\r
+    impl::comp_data(msg_type, left, right, left.channels(), datas.front(), stream);\r
 \r
     for (int i = 1; i < levels; i++)\r
     {\r

diff --git a/modules/gpu/src/cuda/beliefpropagation.cu b/modules/gpu/src/cuda/beliefpropagation.cu
index db0e13a..e7201fe 100644 (file)
--- a/modules/gpu/src/cuda/beliefpropagation.cu
+++ b/modules/gpu/src/cuda/beliefpropagation.cu
@@ -81,26 +81,60 @@ namespace cv { namespace gpu { namespace impl {
 namespace beliefpropagation_gpu\r
 {\r
     template <typename T>\r
-    __global__ void comp_data(uchar* l, uchar* r, size_t step, T* data, size_t data_step, int cols, int rows) \r
+    __global__ void comp_data_gray(const uchar* l, const uchar* r, size_t step, T* data, size_t data_step, int cols, int rows) \r
     {\r
         int x = blockIdx.x * blockDim.x + threadIdx.x;\r
         int y = blockIdx.y * blockDim.y + threadIdx.y;\r
 \r
-        if (y < rows && x < cols)\r
+        if (y > 0 && y < rows - 1 && x > 0 && x < cols - 1)\r
         {\r
-            uchar* ls = l + y * step + x; \r
-            uchar* rs = r + y * step + x; \r
+            const uchar* ls = l + y * step + x; \r
+            const uchar* rs = r + y * step + x; \r
 \r
             T* ds = data + y * data_step + x;\r
             size_t disp_step = data_step * rows;\r
 \r
             for (int disp = 0; disp < cndisp; disp++) \r
             {\r
-                if (x - disp >= 0)\r
+                if (x - disp >= 1)\r
                 {\r
-                    int le = ls[0];\r
-                    int re = rs[-disp];\r
-                    float val = abs(le - re);\r
+                    float val  = abs((int)ls[0] - rs[-disp]);\r
+                    \r
+                    ds[disp * disp_step] = saturate_cast<T>(fmin(cdata_weight * val, cdata_weight * cmax_data_term));\r
+                }\r
+                else\r
+                {\r
+                    ds[disp * disp_step] = saturate_cast<T>(cdata_weight * cmax_data_term);\r
+                }\r
+            }\r
+        }\r
+    }\r
+\r
+    template <typename T>\r
+    __global__ void comp_data_bgr(const uchar* l, const uchar* r, size_t step, T* data, size_t data_step, int cols, int rows) \r
+    {\r
+        int x = blockIdx.x * blockDim.x + threadIdx.x;\r
+        int y = blockIdx.y * blockDim.y + threadIdx.y;\r
+\r
+        if (y > 0 && y < rows - 1 && x > 0 && x < cols - 1)\r
+        {\r
+            const uchar* ls = l + y * step + x * 3; \r
+            const uchar* rs = r + y * step + x * 3; \r
+\r
+            T* ds = data + y * data_step + x;\r
+            size_t disp_step = data_step * rows;\r
+\r
+            for (int disp = 0; disp < cndisp; disp++) \r
+            {\r
+                if (x - disp >= 1)\r
+                {                    \r
+                    const float tr = 0.299f;\r
+                    const float tg = 0.587f;\r
+                    const float tb = 0.114f;\r
+\r
+                    float val  = tb * abs((int)ls[0] - rs[0-disp*3]);\r
+                          val += tg * abs((int)ls[1] - rs[1-disp*3]);\r
+                          val += tr * abs((int)ls[2] - rs[2-disp*3]);\r
                     \r
                     ds[disp * disp_step] = saturate_cast<T>(fmin(cdata_weight * val, cdata_weight * cmax_data_term));\r
                 }\r
@@ -114,10 +148,10 @@ namespace beliefpropagation_gpu
 }\r
 \r
 namespace cv { namespace gpu { namespace impl {\r
-    typedef void (*CompDataFunc)(const DevMem2D& l, const DevMem2D& r, DevMem2D mdata, const cudaStream_t& stream);
+    typedef void (*CompDataFunc)(const DevMem2D& l, const DevMem2D& r, int channels, DevMem2D mdata, const cudaStream_t& stream);
 
     template<typename T>
-    void comp_data_(const DevMem2D& l, const DevMem2D& r, DevMem2D mdata, const cudaStream_t& stream)
+    void comp_data_(const DevMem2D& l, const DevMem2D& r, int channels, DevMem2D mdata, const cudaStream_t& stream)
     {\r
         dim3 threads(32, 8, 1);\r
         dim3 grid(1, 1, 1);\r
@@ -125,13 +159,16 @@ namespace cv { namespace gpu { namespace impl {
         grid.x = divUp(l.cols, threads.x);\r
         grid.y = divUp(l.rows, threads.y);\r
         \r
-        beliefpropagation_gpu::comp_data<T><<<grid, threads, 0, stream>>>(l.ptr, r.ptr, l.step, (T*)mdata.ptr, mdata.step/sizeof(T), l.cols, l.rows);
+        if (channels == 1)\r
+            beliefpropagation_gpu::comp_data_gray<T><<<grid, threads, 0, stream>>>(l.ptr, r.ptr, l.step, (T*)mdata.ptr, mdata.step/sizeof(T), l.cols, l.rows);\r
+        else\r
+            beliefpropagation_gpu::comp_data_bgr<T><<<grid, threads, 0, stream>>>(l.ptr, r.ptr, l.step, (T*)mdata.ptr, mdata.step/sizeof(T), l.cols, l.rows);
         
         if (stream == 0)
             cudaSafeCall( cudaThreadSynchronize() );
     }\r
 \r
-    void comp_data(int msgType, const DevMem2D& l, const DevMem2D& r, DevMem2D mdata, const cudaStream_t& stream)\r
+    void comp_data(int msg_type, const DevMem2D& l, const DevMem2D& r, int channels, DevMem2D mdata, const cudaStream_t& stream)\r
     {\r
         static CompDataFunc tab[8] =
         {
@@ -145,10 +182,10 @@ namespace cv { namespace gpu { namespace impl {
             0                   // user type
         };
 
-        CompDataFunc func = tab[msgType];
+        CompDataFunc func = tab[msg_type];
         if (func == 0)
             cv::gpu::error("Unsupported message type", __FILE__, __LINE__);
-        func(l, r, mdata, stream);\r
+        func(l, r, channels, mdata, stream);\r
     }\r
 }}}\r
 \r
@@ -200,7 +237,7 @@ namespace cv { namespace gpu { namespace impl {
             cudaSafeCall( cudaThreadSynchronize() );
     }\r
 \r
-    void data_step_down(int dst_cols, int dst_rows, int src_rows, int msgType, const DevMem2D& src, DevMem2D dst, const cudaStream_t& stream)\r
+    void data_step_down(int dst_cols, int dst_rows, int src_rows, int msg_type, const DevMem2D& src, DevMem2D dst, const cudaStream_t& stream)\r
     {\r
         static DataStepDownFunc tab[8] =
         {
@@ -214,7 +251,7 @@ namespace cv { namespace gpu { namespace impl {
             0                        // user type
         };
 
-        DataStepDownFunc func = tab[msgType];
+        DataStepDownFunc func = tab[msg_type];
         if (func == 0)
             cv::gpu::error("Unsupported message type", __FILE__, __LINE__);
         func(dst_cols, dst_rows, src_rows, src, dst, stream);\r
@@ -270,7 +307,7 @@ namespace cv { namespace gpu { namespace impl {
             cudaSafeCall( cudaThreadSynchronize() );
     }\r
 \r
-    void level_up_messages(int dst_idx, int dst_cols, int dst_rows, int src_rows, int msgType, DevMem2D* mus, DevMem2D* mds, DevMem2D* mls, DevMem2D* mrs, const cudaStream_t& stream)\r
+    void level_up_messages(int dst_idx, int dst_cols, int dst_rows, int src_rows, int msg_type, DevMem2D* mus, DevMem2D* mds, DevMem2D* mls, DevMem2D* mrs, const cudaStream_t& stream)\r
     {\r
         static LevelUpMessagesFunc tab[8] =
         {
@@ -284,7 +321,7 @@ namespace cv { namespace gpu { namespace impl {
             0                           // user type
         };
 
-        LevelUpMessagesFunc func = tab[msgType];
+        LevelUpMessagesFunc func = tab[msg_type];
         if (func == 0)
             cv::gpu::error("Unsupported message type", __FILE__, __LINE__);
         func(dst_idx, dst_cols, dst_rows, src_rows, mus, mds, mls, mrs, stream);\r
@@ -413,7 +450,7 @@ namespace cv { namespace gpu { namespace impl {
         }
     }\r
 \r
-    void calc_all_iterations(int cols, int rows, int iters, int msgType, DevMem2D& u, DevMem2D& d, DevMem2D& l, DevMem2D& r, const DevMem2D& data, const cudaStream_t& stream)\r
+    void calc_all_iterations(int cols, int rows, int iters, int msg_type, DevMem2D& u, DevMem2D& d, DevMem2D& l, DevMem2D& r, const DevMem2D& data, const cudaStream_t& stream)\r
     {\r
         static CalcAllIterationFunc tab[8] =
         {
@@ -427,7 +464,7 @@ namespace cv { namespace gpu { namespace impl {
             0                             // user type
         };
 
-        CalcAllIterationFunc func = tab[msgType];
+        CalcAllIterationFunc func = tab[msg_type];
         if (func == 0)
             cv::gpu::error("Unsupported message type", __FILE__, __LINE__);
         func(cols, rows, iters, u, d, l, r, data, stream);\r
@@ -496,7 +533,7 @@ namespace cv { namespace gpu { namespace impl {
             cudaSafeCall( cudaThreadSynchronize() );
     }\r
 \r
-    void output(int msgType, const DevMem2D& u, const DevMem2D& d, const DevMem2D& l, const DevMem2D& r, const DevMem2D& data, DevMem2D disp, const cudaStream_t& stream)\r
+    void output(int msg_type, const DevMem2D& u, const DevMem2D& d, const DevMem2D& l, const DevMem2D& r, const DevMem2D& data, DevMem2D disp, const cudaStream_t& stream)\r
     {            \r
         static OutputFunc tab[8] =
         {
@@ -510,7 +547,7 @@ namespace cv { namespace gpu { namespace impl {
             0                // user type
         };
 
-        OutputFunc func = tab[msgType];
+        OutputFunc func = tab[msg_type];
         if (func == 0)
             cv::gpu::error("Unsupported message type", __FILE__, __LINE__);
         func(u, d, l, r, data, disp, stream);\r