From bf5e930468b07a7e6054e7bc734d1f6a6b37596b Mon Sep 17 00:00:00 2001
From: Rok Mandeljc <rok.mandeljc@gmail.com>
Date: Tue, 24 Mar 2015 18:10:22 +0100
Subject: [PATCH] cudastereo: drawColorDisp: enabled CV_32S and CV_32F
 disparity formats

---
 modules/cudastereo/src/cuda/util.cu | 51 +++++++++++++++++++++++++++++
 modules/cudastereo/src/util.cpp     |  6 ++--
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/modules/cudastereo/src/cuda/util.cu b/modules/cudastereo/src/cuda/util.cu
index 6826f90f71..b65c240ee2 100644
--- a/modules/cudastereo/src/cuda/util.cu
+++ b/modules/cudastereo/src/cuda/util.cu
@@ -205,6 +205,29 @@ namespace cv { namespace cuda { namespace device
         }
     }
 
+    __global__ void drawColorDisp(int* disp, size_t disp_step, uchar* out_image, size_t out_step, int width, int height, int ndisp)
+    {
+        const int x = blockIdx.x * blockDim.x + threadIdx.x;
+        const int y = blockIdx.y * blockDim.y + threadIdx.y;
+
+        if(x < width && y < height)
+        {
+            uint *line = (uint*)(out_image + y * out_step);
+            line[x] = cvtPixel(disp[y*disp_step + x], ndisp);
+        }
+    }
+
+    __global__ void drawColorDisp(float* disp, size_t disp_step, uchar* out_image, size_t out_step, int width, int height, int ndisp)
+    {
+        const int x = blockIdx.x * blockDim.x + threadIdx.x;
+        const int y = blockIdx.y * blockDim.y + threadIdx.y;
+
+        if(x < width && y < height)
+        {
+            uint *line = (uint*)(out_image + y * out_step);
+            line[x] = cvtPixel(disp[y*disp_step + x], ndisp);
+        }
+    }
 
     void drawColorDisp_gpu(const PtrStepSzb& src, const PtrStepSzb& dst, int ndisp, const cudaStream_t& stream)
     {
@@ -233,6 +256,34 @@ namespace cv { namespace cuda { namespace device
         if (stream == 0)
             cudaSafeCall( cudaDeviceSynchronize() );
     }
+
+    void drawColorDisp_gpu(const PtrStepSz<int>& src, const PtrStepSzb& dst, int ndisp, const cudaStream_t& stream)
+    {
+        dim3 threads(32, 8, 1);
+        dim3 grid(1, 1, 1);
+        grid.x = divUp(src.cols, threads.x);
+        grid.y = divUp(src.rows, threads.y);
+
+        drawColorDisp<<<grid, threads, 0, stream>>>(src.data, src.step / sizeof(int), dst.data, dst.step, src.cols, src.rows, ndisp);
+        cudaSafeCall( cudaGetLastError() );
+
+        if (stream == 0)
+            cudaSafeCall( cudaDeviceSynchronize() );
+    }
+
+    void drawColorDisp_gpu(const PtrStepSz<float>& src, const PtrStepSzb& dst, int ndisp, const cudaStream_t& stream)
+    {
+        dim3 threads(32, 8, 1);
+        dim3 grid(1, 1, 1);
+        grid.x = divUp(src.cols, threads.x);
+        grid.y = divUp(src.rows, threads.y);
+
+        drawColorDisp<<<grid, threads, 0, stream>>>(src.data, src.step / sizeof(float), dst.data, dst.step, src.cols, src.rows, ndisp);
+        cudaSafeCall( cudaGetLastError() );
+
+        if (stream == 0)
+            cudaSafeCall( cudaDeviceSynchronize() );
+    }
 }}} // namespace cv { namespace cuda { namespace cudev
 
 
diff --git a/modules/cudastereo/src/util.cpp b/modules/cudastereo/src/util.cpp
index a39d28cb2c..09b108ca89 100644
--- a/modules/cudastereo/src/util.cpp
+++ b/modules/cudastereo/src/util.cpp
@@ -92,6 +92,8 @@ namespace cv { namespace cuda { namespace device
 {
     void drawColorDisp_gpu(const PtrStepSzb& src, const PtrStepSzb& dst, int ndisp, const cudaStream_t& stream);
     void drawColorDisp_gpu(const PtrStepSz<short>& src, const PtrStepSzb& dst, int ndisp, const cudaStream_t& stream);
+    void drawColorDisp_gpu(const PtrStepSz<int>& src, const PtrStepSzb& dst, int ndisp, const cudaStream_t& stream);
+    void drawColorDisp_gpu(const PtrStepSz<float>& src, const PtrStepSzb& dst, int ndisp, const cudaStream_t& stream);
 }}}
 
 namespace
@@ -111,11 +113,11 @@ namespace
 void cv::cuda::drawColorDisp(InputArray _src, OutputArray dst, int ndisp, Stream& stream)
 {
     typedef void (*drawColorDisp_caller_t)(const GpuMat& src, OutputArray dst, int ndisp, const cudaStream_t& stream);
-    const drawColorDisp_caller_t drawColorDisp_callers[] = {drawColorDisp_caller<unsigned char>, 0, 0, drawColorDisp_caller<short>, 0, 0, 0, 0};
+    const drawColorDisp_caller_t drawColorDisp_callers[] = {drawColorDisp_caller<unsigned char>, 0, 0, drawColorDisp_caller<short>, drawColorDisp_caller<int>, drawColorDisp_caller<float>, 0, 0};
 
     GpuMat src = _src.getGpuMat();
 
-    CV_Assert( src.type() == CV_8U || src.type() == CV_16S );
+    CV_Assert( src.type() == CV_8U || src.type() == CV_16S || src.type() == CV_32S || src.type() == CV_32F );
 
     drawColorDisp_callers[src.type()](src, dst, ndisp, StreamAccessor::getStream(stream));
 }